#!/usr/bin/env python

""""
Simple implementation of http://arxiv.org/pdf/1502.04623v2.pdf in TensorFlow

Example Usage:
  python train.py --data_file=<TFRecord_file_path> --log_dir=<log_data_path>

Author: Anurag Vempati
"""

import tensorflow as tf
import numpy as np
import os
import sys
import math
import time
from tqdm import tqdm
from config import train_config
from model import DrawModel

tf.flags.DEFINE_string("data_file", "", "")
tf.flags.DEFINE_string("log_dir", "", "")
FLAGS = tf.flags.FLAGS

def export_config(config, output_file):
  """
  Write the configuration parameters into a human readable file.
  :param config: the configuration dictionary
  :param output_file: the output text file
  """
  if not output_file.endswith('.txt'):
      output_file.append('.txt')
  max_key_length = np.amax([len(k) for k in config.keys()])
  with open(output_file, 'w') as f:
      for k in sorted(config.keys()):
          out_string = '{:<{width}}: {}\n'.format(k, config[k], width=max_key_length)
          f.write(out_string)

def load_data(config, data_file):
  print('Loading data from {} ...'.format(data_file))

  # Reads an image from a file, decodes it into a dense tensor, and resizes it
  # to a fixed shape.
  def _parse_function(filename):
    data_fmt = {
      "height": tf.FixedLenFeature((), tf.int64, -1),
      "width": tf.FixedLenFeature((), tf.int64, -1),
      "depth": tf.FixedLenFeature((), tf.int64, -1),
      "image_raw": tf.FixedLenFeature((), tf.string, "")
    }

    parsed_data = tf.parse_single_example(filename, data_fmt)
    image_gray = tf.image.decode_jpeg(parsed_data["image_raw"], channels=1)
    image_converted = tf.image.convert_image_dtype(image_gray, tf.float32)
    image_resized = tf.image.resize_images(image_converted, [config['A'], config['B']])
    image_flattened = tf.reshape(image_resized, [-1])
    return_image = image_flattened
    if not config['draw_with_white']:
      return_image = 1.0 - image_flattened
    return_image = tf.clip_by_value(return_image, 0.0, 0.99)  # for numeric stability during arctanh() operation
    return return_image

  if not os.path.exists(data_file):
    print("Data TFRecord not found")
    sys.exit()

  data_files = tf.data.Dataset.list_files(data_file)
  dataset = data_files.interleave(tf.data.TFRecordDataset, cycle_length=2)
  dataset = dataset.map(map_func=_parse_function, num_parallel_calls=4)
  epoch_counter = tf.data.TFRecordDataset.range(config['n_epochs'])
  dataset = epoch_counter.flat_map(lambda i: tf.data.Dataset.zip(
    (dataset, tf.data.Dataset.from_tensors(i).repeat())))
  dataset = dataset.repeat()
  dataset = dataset.batch(config['batch_size'])
  dataset = dataset.prefetch(buffer_size=config['batch_size'])
  dataset_iterator = dataset.make_one_shot_iterator()
  next_data_batch = dataset_iterator.get_next()

  return next_data_batch

def get_model_and_placeholders(config):
    # create placeholders that we need to feed the required data into the model
    input_pl = tf.placeholder(tf.float32, shape=(config['batch_size'], config['img_size']))
    canvas_pl = tf.placeholder(tf.float32, shape=(config['batch_size'], config['img_size']))
    placeholders = {'input_pl': input_pl,
                    'canvas_pl': canvas_pl}
    return DrawModel, placeholders

def main(config):
  # create unique output directory for this model
  timestamp = str(int(time.time()))
  config['model_dir'] = os.path.abspath(os.path.join(FLAGS.log_dir, 'DRAW' + '_' + timestamp))
  os.makedirs(config['model_dir'])
  print('Logging data to {}'.format(config['model_dir']))

  # Export configuration for the current run
  export_config(config, os.path.join(config['model_dir'], 'config.txt'))

  # load the data
  next_data_batch = load_data(config, FLAGS.data_file)

  # get input placeholders and get the model that we want to train
  draw_model_class, placeholders = get_model_and_placeholders(config)

  # create a training graph, this is the graph we will use to optimize the parameters
  print('Building training graph')
  with tf.name_scope('training'):
    draw_model = draw_model_class(config, placeholders, mode='training', annealing_schedules=config['annealing_schedules'])
    draw_model.build_graph()
    print('created DRAW model with {} parameters'.format(draw_model.n_parameters))

  print('Building valid graph')
  with tf.name_scope('validation'):
    draw_model_valid = draw_model_class(config, placeholders, mode='validation', annealing_schedules=config['annealing_schedules'])
    draw_model_valid.build_graph()
    print('Finished Building valid graphs')

  with tf.Session() as sess:
    # Add the ops to initialize variables.
    init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    # Actually intialize the variables
    sess.run(init_op)

    # Summaries
    training_summaries = tf.summary.merge(tf.get_collection('training_summaries'))
    train_writer = tf.summary.FileWriter(config['model_dir'] + '/summary/train', sess.graph)
    valid_summaries = tf.summary.merge(tf.get_collection('validation_summaries'))
    valid_writer = tf.summary.FileWriter(config['model_dir'] + '/summary/validation', sess.graph)

    # create a saver for writing training checkpoints
    saver = tf.train.Saver(var_list=tf.trainable_variables(), max_to_keep=50)

    # start training
    draw_T = config['T']
#     lowest_test_loss = 1.0e6
    last_saved_epoch = 0  # epoch corresponding to last saved chkpnt
    iteration = 0
    previous_epoch = 0
    epoch = 0
    with tqdm() as pbar:
      while epoch < config['n_epochs']:
        # Next data batch
        previous_epoch = epoch
        xnext, epoch = sess.run(next_data_batch)
        epoch = epoch[0]
        if (previous_epoch > 0 and epoch == 0): break

        step = tf.train.global_step(sess, draw_model.global_step)

        # Hot start
        if config['use_hot_start']:
          crop_fraction = (epoch + 1) * config['crop_fraction_increase_rate']
          if crop_fraction >= 1.0:
            cnext = np.zeros([config['batch_size'], config['img_size']])
            draw_T = config['T']
          else:
            xnext_reshaped = np.copy(xnext)
            xnext_reshaped = xnext_reshaped.reshape((config['batch_size'], config['B'], config['A']))
            start_row = np.random.randint(config['B'] * (1 - crop_fraction))  # , size=config['batch_size'])
            start_col = np.random.randint(config['A'] * (1 - crop_fraction))  # , size=config['batch_size'])
            xnext_reshaped[:, start_row:start_row + int(crop_fraction * config['B']), \
                           start_col:start_col + int(crop_fraction * config['A'])] = 0.0
            cnext = np.reshape(xnext_reshaped, (config['batch_size'], config['img_size']))
            draw_T = max(1, int(config['T'] * crop_fraction))
        else:
          cnext = np.zeros([config['batch_size'], config['img_size']])
          draw_T = config['T']

        # Validate every 100th iteration
        if iteration % 100 == 0:
          valid_feed_dict = draw_model_valid.get_feed_dict(xnext, cnext)
          valid_feed_dict[draw_model_valid.T] = draw_T
          valid_feed_dict[draw_model_valid.global_step] = step
          valid_fetches = {'summaries': valid_summaries,
                           'reconstruction_loss': draw_model_valid.Lx,
                           'latent_loss': draw_model_valid.Lz,
                           'write_loss': draw_model_valid.Lwrite,
                           'intensity_change_loss': draw_model_valid.Lintensity,
                           'movement_loss': draw_model_valid.Lmove,
                           'loss': draw_model_valid.loss}
          valid_out = sess.run(valid_fetches, valid_feed_dict)
          # For saving plot data
          xlog = xnext
          cost = valid_out['loss']
          print("epoch=%d, iter=%d : Lx: %f Lz: %f Lwrite: %f cost: %f" % \
                (epoch, iteration, valid_out['reconstruction_loss'], valid_out['latent_loss'], valid_out['write_loss'], cost))
          valid_writer.add_summary(valid_out['summaries'], global_step=step)
          # save this checkpoint if necessary
          if (epoch - last_saved_epoch + 1) >= config['save_checkpoints_every_epoch']:  # and cost < lowest_test_loss:
            last_saved_epoch = epoch
  #           lowest_test_loss = cost
            saver.save(sess, os.path.join(config['model_dir'], 'drawmodel'), epoch)
        else:
          train_feed_dict = draw_model.get_feed_dict(xnext, cnext)
          train_feed_dict[draw_model.T] = draw_T
          train_fetches = {'summaries': training_summaries,
                           'train_op': draw_model.train_op}
          train_out = sess.run(train_fetches, train_feed_dict)
          if iteration % 100 == 1:
            train_writer.add_summary(train_out['summaries'], global_step=step)

        iteration += 1
        pbar.update(1)

    print('Training finished.')

    # # Logging + Visualization
    log_fetches = {'canvases': draw_model_valid.cs.stack(), 'read_bbs': draw_model_valid.read_bb.stack(), \
                   'write_bbs': draw_model_valid.write_bb.stack(), 'write_times': draw_model_valid.stop_times}
    log_out = sess.run(log_fetches, valid_feed_dict)  # generate some examples
    canvases = np.array(log_out['canvases'])  # T x batch x img_size
    read_bounding_boxes = np.array(log_out['read_bbs'])  # T x batch x 3
    write_bounding_boxes = np.array(log_out['write_bbs'])  # T x batch x 3
    write_times = np.array(log_out['write_times'])  # batch

    log_file = os.path.join(config['model_dir'], "draw_data.npy")
    np.save(log_file, [xlog, canvases, read_bounding_boxes, write_bounding_boxes, write_times, config['draw_with_white']])
    print("Visualization outputs saved in file: %s" % log_file)

    ckpt_file = os.path.join(config['model_dir'], "drawmodel.ckpt")
    print("Model saved in file: %s" % saver.save(sess, ckpt_file))

if __name__ == '__main__':
    main(train_config)


import os

import cv2
import numpy as np

from . import Config
from . import TrackersProvider
from .CalibrationResults import get_current_time
from .QueuesEntries import Point
from .QueuesProvider import QueuesProvider

class Localization(object):
    """Provides localizattion process"""

    objects_count = Config.objects_count
    rotation_matrix1 = None
    rotation_matrix2 = None
    projection_matrix1 = None
    projection_matrix2 = None
    mono_calibration_results = None

    localization_precision = 5  # in millimeters
    last_located_point = [None for _ in range(objects_count)]
    last_located_point_time = [-1 for _ in range(objects_count)]
    time_threshold_skip = 1 / 20
    time_threshold_correspondence = 1 / 10

    @classmethod
    def prepare_projection_matrices(cls, calibration_results1, calibration_results2, stereo_calibration_results):
        '''Prepares projeciton matrices for localization'''
        cls.mono_calibration_results = [calibration_results1, calibration_results2]
        rt = np.append(stereo_calibration_results.rotation_matrix, stereo_calibration_results.translation_vector,
                       axis=1)
        cls.projection_matrix2 = calibration_results2.camera_matrix.dot(rt)
        cls.projection_matrix1 = calibration_results1.camera_matrix.dot(np.eye(3, 4))

    @classmethod
    def get_3d_coordinates(cls, *points):
        '''Undistort points and then compute the estimated 3D position'''
        if any(point is None for point in points) or len(points) != 2:
            return None

        # OpenCv function works with a set of points. In our case we have only one point per camera,
        # therefore we reshape into required format
        points = [np.array(point).reshape(1, 1, 2).astype(float) for point in points]
        points = [cls.get_undistorted_point(point, i) for i, point in enumerate(points)]

        located_points_hom = cv2.triangulatePoints(projMatr1=cls.projection_matrix1, projMatr2=cls.projection_matrix2,
                                                   projPoints1=points[0],
                                                   projPoints2=points[1])
        return cls.convert_from_homogenous(located_points_hom)

    @classmethod
    def get_undistorted_point(cls, point, cam_ind):
        '''Undistort point by using saved calibration data'''
        calib_results = cls.mono_calibration_results[cam_ind]
        undistorted = cv2.undistortPoints(point, calib_results.camera_matrix, calib_results.distortion_coeffs,
                                          P=calib_results.camera_matrix)  # without setting P normalized points would be returned
        return undistorted

    @classmethod
    def save_localization_data(cls):
        '''Saves all localization data of all objects'''
        times = [x[0].timestamp for x in QueuesProvider.LocalizatedPoints3D if x]
        if not times:
            return

        first_point_time = min(times)
        for i, localizated in enumerate(QueuesProvider.LocalizatedPoints3D):
            if not localizated:
                return

            curr_dir = os.path.dirname(os.path.abspath(__file__))
            filename = os.path.join(curr_dir, "localization_data", "{}-{}.txt".format(get_current_time(), i + 1))
            os.makedirs(os.path.dirname(filename), exist_ok=True)

            with open(filename, 'w') as output:
                for point in localizated:
                    point.timestamp -= first_point_time
                    output.write('{}\n'.format(point))

    @classmethod
    def localize_point(cls, object_id):
        '''Localize point in 3D and save it to corresponding queue'''
        points1 = QueuesProvider.TrackedPoints2D[(TrackersProvider.get_tracker_uid(0, object_id))]
        points2 = QueuesProvider.TrackedPoints2D[(TrackersProvider.get_tracker_uid(1, object_id))]

        if len(points1) == 0 or len(points2) == 0:
            return

        point1 = points1[-1]
        point2 = points2[-1]

        if abs(point1.timestamp - point2.timestamp) > cls.time_threshold_correspondence:
            return

        time = (point1.timestamp + point2.timestamp) / 2

        try:
            if time - cls.last_located_point_time[object_id] < cls.time_threshold_skip:
                return
        except IndexError:
            return

        cls.last_located_point_time[object_id] = time

        located_point = Localization.get_3d_coordinates(point1.coordinates,
                                                        point2.coordinates)
        if located_point is None:
            return

        if cls.moved_more_than(cls.last_located_point[object_id], located_point, cls.localization_precision):
            point = Point(located_point, time)
            QueuesProvider.LocalizatedPoints3D[object_id].append(point)
            cls.last_located_point[object_id] = located_point

    @classmethod
    def moved_more_than(cls, old, new, distance):
        return old is None or np.linalg.norm(new - old) > distance

    @classmethod
    def convert_from_homogenous(cls, coords):
        return (coords[:-1] / coords[-1])[:, 0]


"""Provides a class, HP_model, to model the impact of a heat pump on
energy use and cost.
"""
from pprint import pformat
import inspect
from pathlib import Path
import pickle
import time
import gzip

import pandas as pd
import numpy as np

from . import library as lib
from . import elec_cost
from .home_heat_model import HomeHeatModel
from .elec_cost import ElecCostCalc
from .utils import is_null
from . import ui_helper

# --------- Some Constants

# The days in each month
DAYS_IN_MONTH = np.array([
    31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31
])

# The pattern or Lights and appliances other than DHW, Clothes Drying & Cooking.
# This is average power in the month divided average annual power.
LIGHTS_OTHER_PAT = np.array([
    1.13, 1.075, 1.0, 0.925, 0.87, 0.85, 0.87, 0.925, 1.0, 1.075, 1.13, 1.15
])

def make_pattern(esc, life):
    """Makes a numpy array of length (life + 1) containing an escalation pattern
    that starts with a 1.0 in year 1 and escalates at the rate of 'esc' per year.
    """
    pat = np.ones(life - 1) * (1 + esc)
    return np.insert(pat.cumprod(), 0, [0.0, 1.0])

def convert_co2_to_miles_driven(co2_saved):
    """Converts CO2 emissions to a mileage driven
    equivalent for vehicles in the U.S. using EPA
    methodology:  https://www.epa.gov/energy/greenhouse-gases-equivalencies-calculator-calculations-and-references#miles
    """
    pounds_in_metric_ton = 2204.62
    tons_co2_per_gallon = 0.0089
    avg_gas_mileage_us_fleet = 22
    mileage_equivalent = co2_saved / pounds_in_metric_ton / tons_co2_per_gallon * avg_gas_mileage_us_fleet

    return mileage_equivalent

class HP_model:
    """This is the class that orchestrates the heat pump analysis, running the home
    energy model with and without the heat pump.  The class performs the economic
    analysis and stores all results as object attributes.  The main method that
    is run to perform the analysis is "run()".
    """

    # Some of inputs parameters are documented in the home_heat_model.HomeHeatModel class constructor;
    # those inputs are marked as such below.
    def __init__(self,
                 bldg_name,              # name of the building
                 notes,                  # notes about the building
                 city_id,                # see home_heat_model.HomeHeatModel
                 utility,                # The full Pandas Series describing the Electric Utility
                 pce_limit,              # The maximum kWh in a month subsidized by PCE (0 will mean no PCE subsidy)
                 co2_lbs_per_kwh,        # see home_heat_model.HomeHeatModel
                 exist_heat_fuel_id,     # see home_heat_model.HomeHeatModel
                 exist_unit_fuel_cost,   # Cost per physical unit (e.g. gallon, CCF) of the existing space heating fuel.
                 exist_fuel_use,         # Annual existing fuel use for Space Heating and the other end uses identified using that fuel. None if not available.
                 elec_uses,              # 'all' if Annual Electric Use given by user includes lights & appliances, or 'space' if it is just Space Heating use.
                 exist_heat_effic,       # see home_heat_model.HomeHeatModel
                 exist_kwh_per_mmbtu,    # see home_heat_model.HomeHeatModel
                 includes_dhw,           # True if the existing Space Heating Fuel type also is used for DHW.
                 includes_dryer,         # True if the existing Space Heating Fuel type also is used for Clothes Drying.
                 includes_cooking,       # True if the existing Space Heating Fuel type also is used for Cooking.
                 occupant_count,         # Number of occupants using DHW, Clothes Drying, Cooking
                 elec_use_jan,           # The electric use in January, prior to heat pump installation, kWh.
                 elec_use_may,           # The electric use in May, prior to heat pump installation, kWh.
                 hp_model_id,            # see home_heat_model.HomeHeatModel
                 low_temp_cutoff,        # see home_heat_model.HomeHeatModel
                 off_months_chks,        # see home_heat_model.HomeHeatModel, parameter 'off_months' there
                 garage_stall_count,     # see home_heat_model.HomeHeatModel
                 garage_heated_by_hp,    # see home_heat_model.HomeHeatModel
                 bldg_floor_area,        # see home_heat_model.HomeHeatModel
                 indoor_heat_setpoint,   # see home_heat_model.HomeHeatModel
                 insul_level,            # see home_heat_model.HomeHeatModel
                 pct_exposed_to_hp,      # see home_heat_model.HomeHeatModel
                 doors_open_to_adjacent, # see home_heat_model.HomeHeatModel
                 bedroom_temp_tolerance, # see home_heat_model.HomeHeatModel
                 capital_cost,           # Initial cost of the heat pump installation
                 rebate_dol,             # Rebate $ received for heat pump installation
                 pct_financed,           # fraction (0 - 1.0) of heat pump installation cost financed by a loan.
                 loan_term,              # Length of loan in years.
                 loan_interest,          # interest rate of loan, expressed as fraction, i.e. 0.1 for 10%.
                 hp_life,                # life of heat pump in years
                 op_cost_chg,            # operating cost increase associated with heat pump (negative if decrease)
                 sales_tax,              # sales tax, expressed as a fraction (0.05 for 5%/yr) that applies to electricity and fuel costs
                 discount_rate,          # economic discount rate, expressed as a fraction, per year. Nominal, not adjusted for inflation.
                 inflation_rate,         # general inflation rate expressed as a fraction, per year.
                 fuel_esc_rate,          # price escalation rate of fuel used for existing heating system, fraction/year, nominal
                 elec_esc_rate,          # price escalation rate of electricity, fraction/year, nominal
                ):

        # Store all of these input parameters as object attributes.
        args, _, _, values = inspect.getargvalues(inspect.currentframe())
        for arg in args[1:]:
            setattr(self, arg, values[arg])

        # Look up the objects associated with the IDs
        self.city = lib.city_from_id(city_id)
        self.exist_fuel = lib.fuel_from_id(exist_heat_fuel_id)
        self.hp_model = lib.heat_pump_from_id(hp_model_id)

    def __repr__(self):
        """Returns a string with all the object attributes shown.  Text is truncated
        at 1,000 characters for attributes with long representations.
        """
        s = ''
        for attr in self.__dict__:
            val = pformat(self.__dict__[attr])[:1500]
            if len(val)>70:
                s+=f'\n{attr}:\n{val}\n\n'
            else:
                s += f'{attr}: {val}\n'
        return s

    def run(self):
        """This method performs all of the modeling and analysis, storing
        results as object attributes.
        """

        # shortcut for self
        s = self

        # shortcut to existing heating fuel
        fuel = s.exist_fuel

        # holds summary measures for the heat pump project (e.g. seasonal COP,
        # internal rate of return).  Fill out first item: secondary fuel info.
        s.summary = {'fuel_unit': fuel.unit, 'fuel_desc': fuel.desc}

        # Create the home energy simulation object
        sim = HomeHeatModel(
            city_id=s.city_id,
            hp_model_id=s.hp_model_id,
            exist_heat_fuel_id=s.exist_heat_fuel_id,
            exist_heat_effic=s.exist_heat_effic,
            exist_kwh_per_mmbtu=s.exist_kwh_per_mmbtu,
            co2_lbs_per_kwh=s.co2_lbs_per_kwh,
            low_temp_cutoff=s.low_temp_cutoff,
            off_months=s.off_months_chks,
            garage_stall_count=s.garage_stall_count,
            garage_heated_by_hp=s.garage_heated_by_hp,
            bldg_floor_area=s.bldg_floor_area,
            indoor_heat_setpoint=s.indoor_heat_setpoint,
            insul_level=s.insul_level,
            pct_exposed_to_hp=s.pct_exposed_to_hp,
            doors_open_to_adjacent=s.doors_open_to_adjacent,
            bedroom_temp_tolerance=s.bedroom_temp_tolerance,
        )

        # If other end uses use the heating fuel, make an estimate of their annual
        # consumption of that fuel.  This figure is expressed in the physical unit
        # for the fuel type, e.g. gallons of oil.  Save this as an object attribute
        # so it is accessible in other routines.  See Evernote notes on values (AkWarm
        # for DHW and Michael Bluejay for Drying and Cooking).
        is_electric = (s.exist_heat_fuel_id == ui_helper.ELECTRIC_ID)  # True if Electric
        s.fuel_other_uses = s.includes_dhw * 4.23e6 / fuel.dhw_effic
        s.fuel_other_uses += s.includes_dryer * (0.86e6 if is_electric else 2.15e6)
        s.fuel_other_uses += s.includes_cooking * (0.64e6 if is_electric else 0.8e6)
        s.fuel_other_uses *= s.occupant_count / fuel.btus

        # For elecric heat we also need to account for lights and other applicances not
        # itemized above.
        if is_electric:
            # Use the AkWarm Medium Lights/Appliances formula but take 25% off
            # due to efficiency improvements since then.
            s.lights_other_elec = 2086. + 1.20 * s.bldg_floor_area   # kWh in the year
        else:
            s.lights_other_elec = 0.0

        # Match the existing space heating use if it is provided.  Do so by using
        # the UA true up factor.
        if not is_null(s.exist_fuel_use):

            # Remove the energy use from the other end uses that use the fuel, unless
            # this is electric heat and the user indicated that the entered value is
            # just space heating.
            if is_electric and s.elec_uses=='space':
                # user explicitly indicated that the entered annual usage value is
                # just space heating.
                space_fuel_use = s.exist_fuel_use
            else:
                space_fuel_use = s.exist_fuel_use - s.fuel_other_uses - s.lights_other_elec

            sim.no_heat_pump_use = True
            sim.calculate()
            if is_electric:
                # For electric heat, electric use for space heat is in secondary_kwh
                fuel_use1 = sim.annual_results().secondary_kwh
            else:
                fuel_use1 = sim.annual_results().secondary_fuel_units

            # scale the UA linearly to attempt to match the target fuel use
            ua_true_up = space_fuel_use / fuel_use1
            sim.ua_true_up = ua_true_up
            sim.calculate()

            if is_electric:
                # For electric heat, electric use for space heat is in secondary_kwh
                fuel_use2 = sim.annual_results().secondary_kwh
            else:
                fuel_use2 = sim.annual_results().secondary_fuel_units

            # In case it wasn't linear, inter/extrapolate to the final ua_true_up
            slope = (fuel_use2 - fuel_use1)/(ua_true_up - 1.0)
            # print(space_fuel_use, fuel_use1, fuel_use2, ua_true_up)
            ua_true_up = 1.0 + (space_fuel_use - fuel_use1) / slope
            # print(ua_true_up)

        else:
            ua_true_up = 1.0

        # Set the UA true up value into the model and also save it as
        # an attribute of this object so it can be observed.
        sim.ua_true_up = ua_true_up
        s.ua_true_up = ua_true_up

        # Run the base case with no heat pump and record energy results.
        # This model only models the space heating end use.
        sim.no_heat_pump_use = True
        sim.calculate()
        s.df_mo_en_base = sim.monthly_results()
        s.ann_en_base = sim.annual_results()
        # print(s.ann_en_base.secondary_kwh)

        # Run the model with the heat pump and record energy results
        sim.no_heat_pump_use = False
        sim.calculate()
        s.df_mo_en_hp = sim.monthly_results()
        s.ann_en_hp = sim.annual_results()
        s.df_hourly = sim.df_hourly

        # record design heat load
        s.summary['design_heat_load'], s.summary['design_heat_temp'] = sim.design_heat_load()

        # Calculate some summary measures
        s.summary['cop'] = s.ann_en_hp.cop
        s.summary['hp_max_capacity_5F'] = sim.hp_max_capacity_5F()
        s.summary['max_hp_reached'] = sim.max_hp_reached

        # CO2 savings
        s.summary['co2_lbs_saved'] = s.ann_en_base.co2_lbs - s.ann_en_hp.co2_lbs
        s.summary['co2_driving_miles_saved'] = convert_co2_to_miles_driven(s.summary['co2_lbs_saved'])
        s.summary['hp_load_frac'] = s.ann_en_hp.hp_load_mmbtu / (s.ann_en_hp.hp_load_mmbtu + s.ann_en_hp.secondary_load_mmbtu)

        # Create DataFrames that hold monthly energy cost amounts
        # Results are stored as object attributes.
        self.calc_monthly_cash()

        # Create a multi-year Cash Flow DataFrame and summary economic measures.
        # Results are stored as object attributes.
        self.calc_cash_flow()

        # Save a gzipped pickle of this object using Unix time as the file name.
        # make a directory to hold the files
        save_dir = 'hpcalc_runs'
        Path(save_dir).mkdir(exist_ok=True)
        fname = f'{time.time():.2f}.pkl.gz'
        s.file_name = fname
        pickle.dump(self, gzip.open(f'{save_dir}/{fname}', 'wb'))

    def calc_monthly_cash(self):
        """Calculates two DataFrames, s.df_mo_dol_base and s.df_mo_dol_hp, that contain
        the fuel and electricity costs in the base case (no heat pump) scenario and the
        with heat pump scenario.  A number of inputs found as object attributes are used.
        """
        # shortcut to self
        s = self

        # Start the DataFrames, base and w/ heat pump
        # Each starts with just an index column with the month
        # Make shortcut variables as well.
        s.df_mo_dol_base = dfb = s.df_mo_en_base[[]].copy()
        s.df_mo_dol_hp = dfh = s.df_mo_en_base[[]].copy()

        # Determine the base electric use by month.  Approach is different
        # if there is electric heat.
        is_electric_heat = (s.exist_heat_fuel_id == ui_helper.ELECTRIC_ID)
        if not is_electric_heat:
            # Fuel-based space heat.
            # The User supplied a January and a May kWh usage value that should
            # be used for the base case (no heat pump) total electricity use.
            # But, need to come up with a kWh value for every month.  Do that by
            # adjusting the kWh pattern available for this city.
            #
            # Determine the multiplier to adjust to the pattern to the actual.
            pat_use = np.array(s.city.avg_elec_usage)
            mult = (s.elec_use_jan - s.elec_use_may) / (pat_use[0] - pat_use[4])
            pat_use = mult * pat_use
            pat_use += s.elec_use_jan - pat_use[0]

            # The electricity use in the base case
            dfb['elec_kwh'] =  pat_use

            # rough estimate of a base demand: not super critical, as the demand rate
            # structure does not have blocks.  Assume a load factor of 0.4
            dfb['elec_kw'] = dfb.elec_kwh / (DAYS_IN_MONTH * 24.0) / 0.4

        else:
            # Electric Heat Case
            # No Jan and May values are provided.  Instead we have possibly some
            # DHW, clothes drying, and cooking.  Plus, we have base lights/other appliances.
            # And finally we have the Elecric heat making up the base electric usage.

            # First, DHW, Clothes Drying and Cooking.  Assume flat use through year.
            # This is a numpy array because DAYS_IN_MONTH is an array.
            elec_kwh = s.fuel_other_uses / 8760.0 * DAYS_IN_MONTH * 24.0

            # Now lights and other misc. appliances. Some monthly variation, given
            # by LIGHTS_OTHER_PAT.
            elec_kwh += s.lights_other_elec / 8760.0 * LIGHTS_OTHER_PAT * DAYS_IN_MONTH * 24.0

            # For the peak demand of those two categories of use, just assume 40% load factor.
            elec_kw = elec_kwh / (DAYS_IN_MONTH * 24.0) / 0.4

            # Now add in space heating kWh and kW
            elec_kwh += s.df_mo_en_base.total_kwh.values
            elec_kw += s.df_mo_en_base.total_kw.values

            # store results
            dfb['elec_kwh'] =  elec_kwh
            dfb['elec_kw'] =  elec_kw

        # Make an object to calculate electric utility costs
        elec_cost_calc = ElecCostCalc(s.utility, sales_tax=s.sales_tax, pce_limit=s.pce_limit)
        # cost function that will be applied to each row of the cost DataFrame
        cost_func = lambda r: elec_cost_calc.monthly_cost(r.elec_kwh, r.elec_kw)

        dfb['elec_dol'] = dfb.apply(cost_func, axis=1)

        if not is_electric_heat:
            # Now fuel use by month.  Remember that the home heat model only looked at
            # space heating, so we need to add in the fuel use from the other end uses
            # that use this fuel.
            dfb['secondary_fuel_units'] = s.df_mo_en_base.secondary_fuel_units + \
                s.fuel_other_uses / 12.0
            dfb['secondary_fuel_dol'] = dfb.secondary_fuel_units * s.exist_unit_fuel_cost * (1. + s.sales_tax)
        else:
            # Electric Heat, so no secondary fuel
            dfb['secondary_fuel_units'] = 0.0
            dfb['secondary_fuel_dol'] = 0.0

        # Total Electric + space heat
        dfb['total_dol'] =  dfb.elec_dol + dfb.secondary_fuel_dol

        # Now with the heat pump
        # determine extra kWh used in the heat pump scenario. Note, this will
        # be negative numbers if the base case used electric heat.
        extra_kwh = (s.df_mo_en_hp.total_kwh - s.df_mo_en_base.total_kwh).values
        dfh['elec_kwh'] = dfb['elec_kwh'] + extra_kwh
        extra_kw = (s.df_mo_en_hp.total_kw - s.df_mo_en_base.total_kw).values
        dfh['elec_kw'] =  dfb['elec_kw'] + extra_kw
        dfh['elec_dol'] = dfh.apply(cost_func, axis=1)

        # Now fuel, including other end uses using the heating fuel
        if not is_electric_heat:
            dfh['secondary_fuel_units'] = s.df_mo_en_hp.secondary_fuel_units + \
                s.fuel_other_uses / 12.0
            dfh['secondary_fuel_dol'] = dfh.secondary_fuel_units * s.exist_unit_fuel_cost * (1. + s.sales_tax)
        else:
            # Electric Heat, so no secondary fuel
            dfh['secondary_fuel_units'] = 0.0
            dfh['secondary_fuel_dol'] = 0.0

        # Total Electric + space heat
        dfh['total_dol'] =  dfh.elec_dol + dfh.secondary_fuel_dol

    def calc_cash_flow(self):
        """Calculates the cash flow impacts of the installation over the
        life of the heat pump.  Creates a DataFrame, self.df_cash_flow, that shows
        the impacts. In that DataFrame, postive values are benefits and negative
        values are costs.
        Also calculates some summary economic measures that are added to
        the self.summary dictionary.
        """
        s = self   # shortcut variable

        # determine the changes caused by the heat pump on an annual basis.
        # First calculate annual totals for base case and heat pump case and
        # then calculate the change.
        ann_base = s.df_mo_dol_base.sum()
        ann_hp = s.df_mo_dol_hp.sum()
        ann_chg = ann_hp - ann_base
        initial_cost = np.zeros(s.hp_life+1)

        # Am not automatically adding sales tax to the initial cost as the user was
        # supposed to includes sales tax in their input.
        initial_cost[0] = -s.capital_cost * (1 - s.pct_financed) + s.rebate_dol
        loan_pmt = np.pmt(s.loan_interest, s.loan_term, s.capital_cost * s.pct_financed)
        if loan_pmt < -0.01:   # loan payment is negative
            loan_cost = [0.0] + [loan_pmt] * s.loan_term + [0.0] * (s.hp_life -  s.loan_term)
            loan_cost = np.array(loan_cost)
        else:
            loan_cost = 0.0
        op_cost = -s.op_cost_chg * make_pattern(s.inflation_rate, s.hp_life)
        fuel_cost = -ann_chg.secondary_fuel_dol * make_pattern(s.fuel_esc_rate, s.hp_life)
        elec_cost = -ann_chg.elec_dol * make_pattern(s.elec_esc_rate, s.hp_life)
        cash_flow = initial_cost + loan_cost + op_cost + fuel_cost + elec_cost

        # calculate cumulative, discounted cash flow.
        disc_factor = np.ones(s.hp_life) * (1 + s.discount_rate)
        disc_factor = np.insert(disc_factor.cumprod(), 0, 1.0)
        cum_disc_cash_flow = np.cumsum(cash_flow / disc_factor)

        s.df_cash_flow = pd.DataFrame(
            {'initial_cost': initial_cost,
             'loan_cost': loan_cost,
             'op_cost': op_cost,
             'fuel_cost': fuel_cost,
             'elec_cost': elec_cost,
             'cash_flow': cash_flow,
             'cum_disc_cash_flow': cum_disc_cash_flow,
            }
        )
        s.df_cash_flow.index.name = 'year'

        # Calculate IRR and NPV for w/ and w/o PCE.
        s.summary['irr'] = np.irr(s.df_cash_flow.cash_flow)
        s.summary['npv'] = np.npv(s.discount_rate, s.df_cash_flow.cash_flow)

        # Add some summary fuel and electric usage  and unit cost info
        s.summary['fuel_use_base'] = ann_base.secondary_fuel_units
        s.summary['fuel_use_hp'] =  ann_hp.secondary_fuel_units
        s.summary['fuel_use_chg'] = ann_chg.secondary_fuel_units
        if ann_chg.secondary_fuel_units != 0.0:
            s.summary['fuel_price_incremental'] = ann_chg.secondary_fuel_dol / ann_chg.secondary_fuel_units
        else:
            s.summary['fuel_price_incremental'] = np.nan
        s.summary['elec_use_base'] = ann_base.elec_kwh
        s.summary['elec_use_hp'] =  ann_hp.elec_kwh
        s.summary['elec_use_chg'] = ann_chg.elec_kwh
        s.summary['elec_rate_avg_base'] = ann_base.elec_dol / ann_base.elec_kwh
        s.summary['elec_rate_avg_hp'] = ann_hp.elec_dol / ann_hp.elec_kwh
        s.summary['elec_rate_incremental'] = ann_chg.elec_dol / ann_chg.elec_kwh


# Copyright (c) 2021, Oracle and/or its affiliates.  All rights reserved.
# This software is licensed to you under the Universal Permissive License (UPL) 1.0 as shown at
# https://oss.oracle.com/licenses/upl
"""Module contains the MACEst model to estimate prediction intervals."""

import nmslib
import logging
import os
import numpy as np
import scipy
from scipy.optimize import differential_evolution
from scipy.stats import laplace, norm
from scipy.stats._continuous_distns import laplace_gen
from typing import Optional, Union, NamedTuple, Tuple, Dict, Any
from typing_extensions import Protocol, Literal

log = logging.getLogger()

num_threads_available = os.cpu_count()

class _RegressionPointPredictionModel(Protocol):
    """
    Defines a protocol for the type PointPredModel.

    if a model has a predict method then it is the same type as PointPredModel.
    """

    def predict(self, x_star: np.ndarray) -> Any:
        """Return nothing as is only needed to check method exists."""
        pass

class HnswGraphArgs(NamedTuple):
    """Object for passing arguments to the nmslib function."""

    init_kwargs: Dict[str, str] = {"method": "hnsw", "space": "l2"}
    construction_kwargs: Dict[str, int] = {"post": 2, "efConstruction": 1000, "M": 100}
    query_kwargs: Dict[str, int] = {"ef": 1000}

class SearchBounds(NamedTuple):
    """Object for passing the range of allowed MACEst parameters."""

    alpha_bounds: Tuple[float, float] = (0.1, 50.0)
    beta_bounds: Tuple[float, float] = (0.1, 50.0)
    k_bounds: Tuple[int, int] = (5, 20)

class MacestPredIntervalModelParams(NamedTuple):
    """Class container for MACEst model parameters."""

    alpha: float = 1.0
    beta: float = 1.0
    num_neighbours: int = 10

class PrecomputedNeighbourInfo(NamedTuple):
    """Class container for the information about pre-computed nearest neighbours per class."""

    prec_distance_to_nn: Union[Dict[int, np.ndarray], np.ndarray]
    prec_ind_of_nn: Union[Dict[int, np.ndarray], np.ndarray]  # Rhys- only dict for training, where I cache
    # arrays for all values of allowed num neighbours

class ModelWithPredictionInterval:
    """Creates a model which returns a prediction and a confidence interval."""

    def __init__(
        self,
        model: _RegressionPointPredictionModel,
        x_train: np.ndarray,
        train_err: np.ndarray,
        macest_model_params: MacestPredIntervalModelParams = MacestPredIntervalModelParams(),
        error_dist: Literal["normal", "laplace"] = "normal",
        dist_func: Literal["linear", "error_weighted_poly"] = "linear",
        precomputed_neighbour_info: Optional[PrecomputedNeighbourInfo] = None,
        prec_point_preds: Optional[np.ndarray] = None,
        prec_graph: Optional[nmslib.dist.FloatIndex] = None,
        search_method_args: HnswGraphArgs = HnswGraphArgs(),
    ):
        """
        Init.

        :param model: Any model which takes some variables x and returns a point prediction y
        :param x_train: The variables used to train the model
        :param train_err: The error for each training point
        :param num_neighbours: The number of points which define the local neighbourhood
        :param alpha: co-efficient for distance function (hyper-parameter)
        :param beta: The hyper-parameter used in distance function
        :param error_dist: The assumed distribution for the errors
        :param dist_func: The function to convert distance to confidence \
                          (currently linear or error_weighted_poly implemented)
        :param prec_point_preds: The pre-computed model predictions
        :param prec_distance_to_nn: The pre-computed nearest neighbour distances for the calibration and test data
        :param prec_ind_of_nn: The pre-computed nearest neighbour indices for the calibration and test data
        :param prec_graph: The pre-computed graph to use for online hnsw search
        """
        self.model = model
        self.x_train = x_train
        self.train_err = train_err
        self.macest_model_params = macest_model_params
        self._num_neighbours = macest_model_params.num_neighbours
        self._alpha = macest_model_params.alpha
        self._beta = macest_model_params.beta
        self.dist_func = dist_func
        self.error_dist = error_dist
        self.prec_graph = prec_graph
        self.point_preds = prec_point_preds
        self.precomputed_neighbour_info = precomputed_neighbour_info
        if not self.precomputed_neighbour_info:
            self._distance_to_nn = None
            self._ind_of_nn = None
        else:
            self._distance_to_nn = self.precomputed_neighbour_info.prec_distance_to_nn
            self._ind_of_nn = self.precomputed_neighbour_info.prec_ind_of_nn
        self.search_method_args = search_method_args
        self._check_consistent_search_method_args()
        self._check_data_consistent_with_search_args()

    def predict(self, x_star: np.ndarray) -> np.ndarray:
        """
        Return a point prediction for x_star.

        :param x_star: The position for which we would like to predict

        :return: pred_star : The point prediction for x_star
        """
        pred_star = self.model.predict(x_star)
        return pred_star

    def build_graph(self) -> nmslib.dist.FloatIndex:
        """
        Build the  Hierarchical Navigable Small World (hnsw) index graph.

        :return: A queryable HNSW graph
        """
        graph = nmslib.init(**self.search_method_args.init_kwargs)
        graph.addDataPointBatch(self.x_train)
        graph.createIndex(self.search_method_args.construction_kwargs)
        graph.setQueryTimeParams(self.search_method_args.query_kwargs)

        return graph

    def calc_nn_dist(self, x_star: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate the distant to a set of k nearest neighbours.

        :param x_star: The position for which we would like to predict

        :return: The distance to k nearest neighbours and the indices of the k closest neighbours
        """
        if self.prec_graph is None:
            self.prec_graph = self.build_graph()

        neighbours = np.array(
            self.prec_graph.knnQueryBatch(
                x_star, k=self._num_neighbours, num_threads=num_threads_available
            )
        )
        dist = neighbours[:, 1, :]
        ind = neighbours[:, 0, :].astype(int)
        return dist, ind

    def calc_linear_dist_func(self, x_star: np.ndarray) -> np.ndarray:
        """
        Calculate the linear sum of average distance to neighbours and average per neighbour error.

        :param x_star: The position for which we would like to predict

        :return: the sum of average distance to neighbours and average per neighbour error for x_star
        """
        if self._distance_to_nn is not None:
            local_distance = self._distance_to_nn
            if self._ind_of_nn is None:
                raise ValueError("_ind_of_nn has not been cached during training")
            ind = self._ind_of_nn
        else:
            local_distance, ind = self.calc_nn_dist(x_star)
        if isinstance(local_distance, np.ndarray):
            dist = self._alpha * np.average(
                local_distance, weights=np.arange(local_distance.shape[1], 0, -1), axis=1,
            )
        else:
            raise ValueError('Need to remove pre-cached training neighbour data from training')
        if isinstance(ind, np.ndarray):
            error = self._beta * np.average(
                abs(self.train_err[ind.astype(int)]),
                weights=1.0 / (1 + local_distance),
                axis=1,)
        else:
            raise ValueError('Need to remove pre-cached training neighbour data from training')

        return dist + error

    def calc_error_weighted_dist(self, x_star: np.ndarray) -> np.ndarray:
        """
        Calculate average distance to neighbours weighted by the per neighbour prediction error.

        :param x_star: The position for which we would like to predict

        :return: the error weighted distance from x_star point to it's neighbours
        """
        if self._distance_to_nn is not None:
            local_distance = self._distance_to_nn
            if self._ind_of_nn is None:
                raise ValueError("_ind_of_nn has not been cached during training")
            ind = self._ind_of_nn
        else:
            local_distance, ind = self.calc_nn_dist(x_star)

        if isinstance(ind, np.ndarray):
            train_error = self.train_err[ind.astype(int)]
        else:
            raise ValueError('Need to remove pre-cached training neighbour data from training')
        if isinstance(local_distance, np.ndarray):
            error_weighted_dist = np.average(
                local_distance * abs(train_error),
                weights=1.0 / (1 + local_distance),
                axis=1,
            )
        else:
            raise ValueError('Need to remove pre-cached training neighbour data from training')

        error_weighted_poly = self._alpha * error_weighted_dist ** self._beta
        return error_weighted_poly

    def std_on_y_star(self, x_star: np.ndarray) -> np.ndarray:
        """
        Return the predicted variance for x_star.

        :param x_star: The position for which we would like to predict

        :return: sigma: The standard deviation for the prediction at x_star
        """
        if self.dist_func == "error_weighted_poly":
            dist = self.calc_error_weighted_dist(x_star)
        elif self.dist_func == "linear":
            dist = self.calc_linear_dist_func(x_star)
        else:
            raise ValueError(f"Unknown distance function: {self.dist_func}")
        sigma = dist
        return sigma

    def laplace_scale_on_y_star(self, x_star: np.ndarray) -> np.ndarray:
        """
        Return the predicted laplacian variance for x_star.

        :param x_star: The position for which we would like to predict

        :return: sigma: The laplacian scaler for the prediction at x_star
        """
        if self.dist_func == "error_weighted_poly":
            dist = self.calc_error_weighted_dist(x_star)
        elif self.dist_func == "linear":
            dist = self.calc_linear_dist_func(x_star)
        else:
            raise ValueError(f"Unknown distance function: {self.dist_func}")
        sigma = dist
        return sigma

    def _distribution(self, x_star: np.ndarray) -> laplace_gen:
        """
        Return the distribution that we will predict from.

        :return:
        """
        if self.point_preds is not None:
            point_preds = self.point_preds
        else:
            point_preds = self.predict(x_star,)
        if self.error_dist == "normal":
            scale = self.std_on_y_star(x_star,)
            dist = norm(loc=point_preds, scale=scale)
        elif self.error_dist == "laplace":
            scale = self.laplace_scale_on_y_star(x_star,)
            dist = laplace(loc=point_preds, scale=scale)
        else:
            raise ValueError(f"Unknown distance function: {self.dist_func}")
        return dist

    def predict_interval(
        self, x_star: np.ndarray, conf_level: Union[np.ndarray, int, float] = 90,
    ) -> np.ndarray:
        """
        Predict the upper and lower prediction interval bounds for a given confidence level.

        :param x_star: The position for which we would like to predict
        :param conf_level:

        :return: The confidence bounds for each x_star for each confidence level
        """
        dist = self._distribution(x_star)
        lower_perc = (100 - conf_level) / 2
        upper_perc = 100 - lower_perc

        lower_vec = 0.01 * np.ones((x_star.shape[0], len([conf_level]))) * lower_perc
        upper_vec = 0.01 * np.ones((x_star.shape[0], len([conf_level]))) * upper_perc
        return np.array([dist.ppf(lower_vec.T), dist.ppf(upper_vec.T)]).T

    def calculate_prediction_interval_width(
        self, x_star: np.ndarray, conf_level: Union[np.ndarray, int, float] = 90,
    ) -> np.ndarray:
        """
        Calculate the absolute width of a prediction interval for a given confidence level.

        :param x_star: The position for which we would like to predict
        :param conf_level:

        :return: the absolute width of a prediction interval for each x_star for each confidence level
        """
        intervals = self.predict_interval(x_star, conf_level)
        return np.diff(intervals)

    def sample_prediction(
        self, x_star: np.ndarray, nsamples: int = 10 ** 3
    ) -> np.ndarray:
        """
        Draw samples from any predicted distribution to get a distribution of predictions.

        :param x_star: The position in feature space for which we would like to predict
        :param nsamples: The number of samples to draw from the distribution

        :return: Samples from the predicted distribution
        """
        dist = self._distribution(x_star)
        return dist.rvs(size=(nsamples, x_star.shape[0])).T

    def fit(
        self,
        x_cal: np.ndarray,
        y_cal: np.ndarray,
        param_range: SearchBounds = SearchBounds(),
        optimiser_args: Optional[Dict[Any, Any]] = None,
    ) -> None:
        """
        Fit MACEst model using the calibration data.

        :param x_cal: Calibration data
        :param y_cal: Target values
        :param param_range: The bounds within which to search for MACEst parameters
        :param optimiser_args: Any arguments for the optimiser (see scipy.optimize)

        :return: None
        """
        if optimiser_args is None:
            optimiser_args = {}

        train_helper = _TrainingHelper(self, x_cal, y_cal, param_range)
        train_helper.fit(optimiser_args=optimiser_args)

    def _check_consistent_search_method_args(self) -> None:
        init_args = self.search_method_args.init_kwargs
        index = nmslib.init(**init_args)

        if 'space' not in list(init_args.keys()):
            raise ValueError('You must pass a space in your search method init args')

        space = init_args['space']
        if space[-6:] == 'sparse':
            sparse_metric = True
        else:
            sparse_metric = False

        data_type = index.dataType
        if data_type == nmslib.DataType.SPARSE_VECTOR:
            sparse_data = True
        else:
            sparse_data = False

        if sparse_metric != sparse_data:
            raise ValueError(
                f'Data type and space are not compatible, your space is {space} '
                f'and search data type is data_type nmslib.{data_type}')

    def _check_data_consistent_with_search_args(self) -> None:
        init_args = self.search_method_args.init_kwargs

        space = init_args['space']
        if space[-6:] == 'sparse':
            sparse_metric = True
        else:
            sparse_metric = False

        training_data_type = type(self.x_train)

        if training_data_type == scipy.sparse.csr.csr_matrix:
            sparse_data = True
        else:
            sparse_data = False

        if sparse_metric != sparse_data:
            raise ValueError(
                f'Training data type and space are not compatible, your space is {space} '
                f'and training data type is {training_data_type}')

class _TrainingHelper(object):
    def __init__(
        self,
        init_conf_model: ModelWithPredictionInterval,
        x_cal: np.ndarray,
        y_cal: np.ndarray,
        param_range: SearchBounds = SearchBounds(),
    ):
        """
        Init.

        :param init_conf_model: an initialised ModelWithConfidence object that we want to fit
        :param x_cal: The X variables that we will use to calibrate the confidence predictions
        :param y_cal: The target variables that we will use to calibrate the confidence predictions
        :param param_range: The bounds on the hyper-parameter space we want to search
        """
        self.model = init_conf_model
        self.x_cal = x_cal
        self.y_cal = y_cal
        self.param_range = param_range
        self.prec_graph = self.model.build_graph()
        self.model.prec_graph = self.prec_graph
        self.prec_dist, self.prec_ind = self._prec_neighbours()
        self.model.point_preds = self.model.predict(self.x_cal)

    def _prec_neighbours(self) -> Tuple[Dict[int, np.ndarray], Dict[int, np.ndarray]]:
        """
        Pre-compute the nearest neighbours and their distances.

        :return:
        """
        min_nbrs = self.param_range[2][0]
        max_nbrs = self.param_range[2][1]
        num_nbrs = np.arange(min_nbrs, max_nbrs + 0.1, 1)
        x_cal_len_array = np.arange(len(self.x_cal))

        dist_dict = {}
        ind_dict = {}

        max_neighbours = np.array(
            self.prec_graph.knnQueryBatch(
                self.x_cal, k=int(max_nbrs), num_threads=num_threads_available
            )
        )

        max_dist = max_neighbours[x_cal_len_array, 1]
        max_ind = max_neighbours[x_cal_len_array, 0]

        for k in num_nbrs:
            dist = max_dist[x_cal_len_array, 0: int(k)]
            ind = max_ind[x_cal_len_array, 0: int(k)]

            dist_dict[k] = dist
            ind_dict[k] = ind

        return dist_dict, ind_dict

    def set_macest_model_params(self) -> MacestPredIntervalModelParams:
        """
        Return MACEst parameter values.

        :return:
        """
        params = MacestPredIntervalModelParams(
            num_neighbours=self.model._num_neighbours,
            alpha=self.model._alpha,
            beta=self.model._beta,
        )
        self.model.macest_model_params = params
        return params

    def loss_func(self, params: MacestPredIntervalModelParams) -> float:
        """
        Calculate the loss for a given set of parameters, this will then be optimised when fit is called.

        :param params: A tuple containing the model hyper-paramters
        :return:
        """
        self.model._alpha, self.model._beta, self.model._num_neighbours = params

        self.model._num_neighbours = int(np.round(self.model._num_neighbours))

        self.model.prec_graph = self.prec_graph
        self.model._distance_to_nn = self.prec_dist[self.model._num_neighbours]
        self.model._ind_of_nn = self.prec_ind[self.model._num_neighbours]

        return picp_loss(self.model, self.x_cal, self.y_cal)

    def fit(
        self,
        optimiser: Literal["de"] = "de",
        optimiser_args: Optional[Dict[Any, Any]] = None,
    ) -> ModelWithPredictionInterval:
        """
        Fit MACEst parameters.

        :param optimiser: The optimisation method
        :param optimiser_args: Any arguments for the optimisation strategy
        :return: A ModelWithConfidence object with the hyper-parameters that minimises the loss function
        """
        if optimiser == "de":
            result = differential_evolution(
                self.loss_func, self.param_range, **optimiser_args
            )
        else:
            raise ValueError(
                "The only optimisation method currently implemented is differential evolution"
            )

        log.info(f"min_loss = {result.fun}")

        alpha, beta, k = result.x
        k = int(np.round(k, 0))
        log.info(f" best_alpha: {alpha}")
        log.info(f" best_beta: {beta}")
        log.info(f" best_k: {k}")

        self.model._alpha = alpha
        self.model._beta = beta
        self.model._num_neighbours = int(np.round(k))

        self.model.macest_model_params = self.set_macest_model_params()

        self.model._distance_to_nn = None
        self.model._ind_of_nn = None
        self.model.point_preds = None

        return self.model

def picp_loss(
    interval_model: ModelWithPredictionInterval, x_test: np.ndarray, y_true: np.ndarray
) -> float:
    """
    Calculate the difference between the desired confidence level and the \
    prediction_interval_coverage_probability for several intervals.

    :param interval_model: Some model which makes predictions with a standard deviation
    :param x_test: The variables for which we would like to use to predict a distribution
    :param y_true: The True target value

    :return: The loss score
    """
    levels = np.array((90, 70, 50, 30, 10))

    intervals = interval_model.predict_interval(x_test, conf_level=levels,).T

    lower = intervals[0]
    upper = intervals[1]
    loss = 0

    for i in range(len(levels)):
        loss += abs(
            levels[i]
            - (
                100
                * len(
                    np.where(np.logical_and(y_true >= lower[i], y_true <= upper[i]))[0]
                )
                / len(y_true)
            )
        )
    return loss / len(levels)


"""
This is the code associated to the paper "Low Dimensional State Representation Learning with Robotics Priors in Continuous Action Spaces"
Botteghi N., et al, published in the International Conference of Intelligent Systems and Robots (IROS), September 2021.

The StateRepresentation class includes the encoder neural network for learning a low-dimensional state representation
from high-dimesional observations (lidar data points + RGB camera images). The encoder is trained with a new set of
robotics priors tailored for continuous state and action spaces.
"""

import tensorflow as tf
# Hide some depreacation warnings and disable eager execution
tf.logging.set_verbosity(tf.logging.ERROR)
import os
import numpy as np
import sonnet as snt
from utils import load_pickle, reshape_observation, save_pickle
from tqdm import tqdm

class ReplayBuffer:
    def __init__(self, obs_dim, act_dim, size):
        self.obs_dim = obs_dim
        self.act_dim = act_dim
        self.size = size

        self.obs_buf = np.zeros([int(size), int(obs_dim)], dtype=np.float32)
        self.acts_buf = np.zeros([int(size), int(act_dim)], dtype=np.float32)
        self.rews_buf = np.zeros(int(size), dtype=np.float32)
        self.done_buf = np.zeros(int(size), dtype=np.float32)
        self.sample_nr_buf = np.zeros(int(size), dtype=np.float32)
        self.ptr, self.size, self.max_size = 0, 0, int(size)

    def store(self, obs, act, rew, done, sample_nr):
        self.obs_buf[self.ptr] = obs
        self.acts_buf[self.ptr] = act
        self.rews_buf[self.ptr] = rew
        self.done_buf[self.ptr] = done
        self.sample_nr_buf[self.ptr] = sample_nr
        self.ptr = (self.ptr + 1) % self.max_size  # replace oldest entry from memory
        self.size = min(self.size + 1, self.max_size)

    # def _get_act_seq

    def sample_batch(self, batch_size=32):
        idxs = np.random.randint(0, self.size - 1, size=batch_size)
        idxs = idxs[self.done_buf[idxs] != 1.]  # remove the last samples of the sequence

        obs_dict = dict(obs=self.obs_buf[idxs],
                        acts=self.acts_buf[idxs],
                        rews=self.rews_buf[idxs - 1],
                        done=self.done_buf[idxs],
                        sample_nr=self.sample_nr_buf[idxs])
        next_obs_dict = dict(obs=self.obs_buf[idxs + 1],
                             acts=self.acts_buf[idxs + 1],
                             rews=self.rews_buf[idxs],
                             done=self.done_buf[idxs + 1],
                             sample_nr=self.sample_nr_buf[idxs + 1])
        return obs_dict, next_obs_dict

    def get_all_samples(self):
        return dict(obs=self.obs_buf[:self.size],
                    acts=self.acts_buf[:self.size],
                    rews=self.rews_buf[:self.size],
                    done=self.done_buf[:self.size],
                    sample_nr=self.sample_nr_buf[:self.size])

    def remove_all(self):
        self.__init__(self.obs_dim, self.act_dim, self.size)

class StateRepresentation(object):
    """This class takes care of learning the state representation"""

    def __init__(self, obs_dim, state_dim=5, act_dim=2, batch_size=256, learning_rate=5e-4,
                 alpha=2, beta=10, seed=1, continuelearning=False, usingros=False):

        folder = "training_results"
        filename = 'SRLnetwork'
        if usingros:
            import rospkg
            self.model_path = os.path.join(rospkg.RosPack().get_path("rosbot_srl"), folder, "srl", filename + ".ckpt")
        else:
            self.model_path = os.path.join("..", folder, "srl", filename + ".ckpt")

        self.obs_dim = obs_dim   # size [32, 24, 40] with 32x24x3 --> size RGB image and 40 --> size of lidar array
        self.obs_count = obs_dim[0] * obs_dim[1] * 3 + obs_dim[2]  #  calculate the flattened dimension of the observation vector
        self.state_dim = state_dim
        self.act_dim = act_dim   # linear and angular velocity of the robot
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.alpha = alpha   # priors coefficient (see paper for more details)
        self.beta = beta   # priors coefficient (see paper for more details)
        self.continuelearning = continuelearning

        self.memory = ReplayBuffer(self.obs_count, act_dim=act_dim, size=3e4)

        tf.random.set_random_seed(seed=seed)
        np.random.seed(seed=seed)

        # DEFINE THE FUNCTIONS REQUIRED FOR TRAINING
        self.graph = tf.Graph()
        with self.graph.as_default():
            # Define the placeholders
            self.obs_1 = tf.placeholder(tf.float32, shape=[None, self.obs_count], name='observation_1')
            self.obs_2 = tf.placeholder(tf.float32, shape=[None, self.obs_count], name='observation_2')
            self.obs_3 = tf.placeholder(tf.float32, shape=[None, self.obs_count], name='observation_3')
            self.obs_4 = tf.placeholder(tf.float32, shape=[None, self.obs_count], name='observation_4')

            self.act_1 = tf.placeholder(tf.float32, shape=[None, 2], name='action_1')
            self.act_2 = tf.placeholder(tf.float32, shape=[None, 2], name='action_2')

            self.is_training = tf.placeholder(tf.bool, shape=[], name="train_cond")

            # Define the neural network architecture and its output
            self.nn = snt.Module(self.SRLencoder, name='SRL_Network')
            self.state_1 = self.nn(self.obs_1, self.is_training)
            self.state_2 = self.nn(self.obs_2, self.is_training)

            self.state_delta = self.state_2 - self.state_1

            self.state_shuff = self.nn(self.obs_3, self.is_training)
            self.state_4 = self.nn(self.obs_4, self.is_training)
            self.state_delt_shuff = self.state_4 - self.state_shuff

            # define losses (i.e. the robotics priors) and optimizer (ADAM)

            self.temp_coh_loss = self.temporal_coherence_prior(self.state_delta, self.act_1,alpha=self.alpha)

            self.caus_loss = self.causality_prior(self.state_1, self.state_shuff, self.act_1,self.act_2, self.beta)

            self.prop_loss = self.proportionality_prior(self.state_delta, self.state_delt_shuff, self.act_1, self.act_2,
                                                        self.beta)

            self.repeat_loss = self.repeatability_prior(self.state_1, self.state_shuff, self.state_delta,
                                                        self.state_delt_shuff, self.act_1, self.act_2, self.beta)

            graph_regularizers = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
            total_regularization_loss = tf.reduce_sum(graph_regularizers)

            self.losses = [0, 1*self.temp_coh_loss, 2*self.caus_loss, 1*self.prop_loss, 1*self.repeat_loss,
                           total_regularization_loss]
            self.loss = tf.reduce_sum(self.losses)
            self.losses[0] = tf.reduce_sum(self.losses[:])

            optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate)
            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.train_op = optimizer.minimize(self.loss)

            self.saver = tf.train.Saver()
            self.init = tf.global_variables_initializer()

        # Initialize the session
        self.sess = tf.Session(graph=self.graph)
        # Restore session if desired
        if continuelearning:
            print('Loading SR model from memory')
            self.load_model()
        else:
            print("Starting SR model from scratch!")
            self.sess.run(self.init)

    def SRLencoder(self, observations, is_training, l2_reg=0.001, batch_norm=False):

        """
        The camera image and the lidar data points are separated from the flattened observation vector via snt.SliceByDim()
        in camera_inputs and laser_inputs. camera_inputs are reshaped into [32, 24, 3] and fed to conv2d layers,
        while laser_inputs are reshaped into [40, 1] and fed to a conv1d layer. Noise is added to the output during training
        for preventing equal state predictions (for more details we reference to "Learning State Representations with Robotics Priors"
        Jonschkowski R. et al).

         image               laser
           |                   |
         conv2d              conv1d
           |                   |
         conv2d                |
           |                   |
        flatten             flatten
           |                   |
         dense               dense
           |                   |
           --------merge--------
                     |
                   dense
                     |
                   dense
                     |
                   dense
                     |
                   state

        :param observations:
        :param is_training:
        :param l2_reg:
        :param batch_norm:
        :return state:
        """

        regularizers = {"w": tf.contrib.layers.l2_regularizer(scale=l2_reg)}
        initializers = {"w": tf.keras.initializers.he_normal()}

        # Camera branch
        camera_inputs = snt.SliceByDim(dims=[1], begin=[0], size=[self.obs_dim[0] * self.obs_dim[1] * 3])(observations)
        camera_inputs = tf.reshape(camera_inputs, [-1, self.obs_dim[0], self.obs_dim[1], 3])

        camera_conv1 = tf.layers.conv2d(camera_inputs, filters=32, kernel_size=3, strides=1, padding='valid',
                                        activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                        kernel_initializer=initializers["w"])

        if batch_norm:
            camera_conv1 = tf.layers.batch_normalization(camera_conv1, training=is_training)

        camera_conv2 = tf.layers.conv2d(camera_conv1, filters=64, kernel_size=3, strides=1, padding='valid',
                                        activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                        kernel_initializer=initializers["w"])
        if batch_norm:
            camera_conv2 = tf.layers.batch_normalization(camera_conv2, training=is_training)

        camera_flatten = snt.BatchFlatten()(camera_conv2)

        camera_dense1 = tf.layers.dense(camera_flatten, 64, activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                        kernel_initializer=initializers["w"])

        # Laser branch
        laser_inputs = snt.SliceByDim(dims=[1], begin=[self.obs_dim[0] * self.obs_dim[1] * 3], size=[self.obs_dim[2]])(
            observations)
        laser_inputs = tf.reshape(laser_inputs, [-1, self.obs_dim[2], 1])

        laser_conv1 = tf.layers.conv1d(laser_inputs, filters=32, kernel_size=3, strides=1, padding='valid',
                                       activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                       kernel_initializer=initializers["w"])

        if batch_norm:
            laser_conv1 = tf.layers.batch_normalization(laser_conv1, training=is_training)

        laser_flatten = snt.BatchFlatten()(laser_conv1)

        laser_dense1 = tf.layers.dense(laser_flatten, 64, activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                       kernel_initializer=initializers["w"])

        # Merged layers
        merge_input = tf.concat(values=[camera_dense1, laser_dense1], axis=1)
        merge_dense1 = tf.layers.dense(merge_input, 64, activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                       kernel_initializer=initializers["w"])
        merge_dense2 = tf.layers.dense(merge_dense1, 32, activation=tf.nn.relu, kernel_regularizer=regularizers["w"],
                                       kernel_initializer=initializers["w"])

        state = tf.layers.dense(merge_dense2, self.state_dim, activation=None,
                                      kernel_regularizer=regularizers["w"], kernel_initializer=tf.initializers.zeros)

        merge_noise = lambda x: tf.cond(self.is_training, lambda: x + tf.random_normal(shape=tf.shape(x), stddev=1e-6),
                                        lambda: x)
        return merge_noise(state)

    def remember(self, observation, action, reward, done, sample_number):
        self.memory.store(observation, action, reward, done, sample_number)

    def learn(self):
        with tqdm(total=self.memory.size) as pbar:
            t_loss = 0
            losses = np.zeros(6)
            nr_batches = self.memory.size // self.batch_size
            for _ in range(nr_batches):
                b1 = self.memory.sample_batch(self.batch_size)
                b2 = self.memory.sample_batch(self.batch_size)

                l = min(len(b1[0].get('obs')), len(b2[0].get('obs')))

                feed_dict = {self.obs_1: b1[0].get('obs')[:l], self.obs_2: b1[1].get('obs')[:l],
                             self.obs_3: b2[0].get('obs')[:l], self.obs_4: b2[1].get('obs')[:l],
                             self.act_1: b1[0].get('acts')[:l],
                             self.act_2: b2[0].get('acts')[:l],
                             self.is_training: True}

                _, loss, l = self.sess.run([self.train_op, self.loss, self.losses], feed_dict=feed_dict)
                t_loss += loss
                losses += l
                pbar.update(self.batch_size)

        return t_loss / nr_batches, losses / nr_batches

    def predict(self, observation):
        feed_dict = {self.obs_1: np.reshape(observation, (1, -1)),
                     self.is_training: False}
        state = self.sess.run(self.state_1, feed_dict=feed_dict)
        return state

    def predict_all(self, observations, batch_size=512):
        observations = observations.reshape(len(observations), -1)
        states = np.ndarray([len(observations), self.state_dim])
        num_batches = int(np.trunc((observations.shape[0] / batch_size)))

        for i in range(num_batches):
            states[i * batch_size: (i + 1) * (batch_size)] = self.sess.run(self.state_1, feed_dict={
                self.obs_1: observations[i * batch_size: (i + 1) * (batch_size)], self.is_training: False})
        states[num_batches * batch_size:] = self.sess.run(self.state_1, feed_dict={
            self.obs_1: observations[num_batches * batch_size:], self.is_training: False})
        return states

    def get_memory_states(self):
        states = []
        mem = self.memory.get_all_samples().get('obs')

        states = self.predict_all(mem)
        return states, self.memory.get_all_samples()

    def save_model(self):
        print("SRL Network storing model..........")
        return self.saver.save(self.sess, self.model_path)

    def load_model(self):
        print("SRL Network restoring data...........")
        self.saver.restore(self.sess, self.model_path)

    # Each prior is defined in a function

    def temporal_coherence_prior(self, s_d, a1, alpha=2):
        """
        E[||s_t+1 - s_t||_2 * e^(-alpha * ||a||_2)]

        :param s_d: s_t+1 - s_t
        :param a1:  action connecting s_t and s_t+1
        :param alpha:
        :return: expectation of the temporal coherence loss
        """

        return tf.reduce_mean(tf.math.exp(-alpha * tf.norm(a1, ord=2, axis=1)) * tf.norm(s_d, ord=2, axis=1) ** 2)

    def causality_prior(self, s1, s2, a1, a2, beta):
        """
        E[e^(-||s1-s2||^2) * e^(-beta*||a1-a2||^2)]

        :param s1: state s_t1
        :param s2: state s_t2
        :param a1: action connecting s_t1 and s_t1+1
        :param a2:  action connecting s_t2 and s_t2+1
        :param beta: weighting factor
        :return: expectation of the causality loss
        """
        closs1 = tf.math.exp(-tf.norm(s1 - s2, ord=2, axis=1) ** 2)
        closs2 = tf.math.exp(-beta*tf.norm(a2 - a1, ord=2, axis=1) ** 2)

        return tf.reduce_mean(closs1 * closs2)

    def proportionality_prior(self, sd1, sd2, a1, a2, beta):
        """
        E[(||s_t2+1 - s_t2||_2 - ||s_t1+1 - s_t1||_2)^2 * e^(-beta * ||a1 - a2||_2)^2]

        :param sd1: s_t1+1 - s_t1
        :param sd2: s_t2+1 - s_t2
        :param a1: action connecting s_t1 and s_t1+1
        :param a2: action connecting s_t2 and s_t2+1
        :param beta: weighting factor
        :return: expectation of the proportinality loss
        """
        ploss1 = (tf.norm(sd2, ord=2, axis=1) - tf.norm(sd1, ord=2, axis=1)) ** 2
        ploss2 = tf.math.exp(-beta * tf.norm(a1 - a2, ord=2, axis=1) ** 2)

        return tf.reduce_mean(ploss1 * ploss2)

    def repeatability_prior(self, s1, s2, s_d1, s_d2, a1, a2, beta):
        """
        E[||(s_t2+1 - s_t2) - (s_t1+1 - s_t1)||_2)^2 * e^(-||s1-s2||^2) * e^(-beta * ||a1 - a2||_2)^2]

        :param s1: state s_t1
        :param s2: state s_t2
        :param s_d1: s_t1+1 - s_t1
        :param s_d2: s_t2+1 - s_t2
        :param a1: action connecting s_t1 and s_t1+1
        :param a2: action connecting s_t2 and s_t2+1
        :param beta: weighting factor
        :return: expectation of the repeatability loss
        """
        rloss1 = tf.math.exp(-tf.norm(s1 - s2, ord=2, axis=1) ** 2)
        rloss2 = tf.norm(s_d2 - s_d1, ord=2, axis=1) ** 2
        rloss3 = tf.math.exp( -beta*tf.norm(a1 - a2, ord=2, axis=1) ** 2)

        return tf.reduce_mean(rloss1 * rloss2 * rloss3)

if __name__ == '__main__':
    from Logger import Logger
    import Plotter as plotter

    # Load saved observations from memory (samples collected by random exploring the large 4 walls environment --> see our paper)
    folder = 'training_data/observations_4walls_large.pkl'

    data = load_pickle(folder)
    print('Loaded {} data points'.format(len(data)))

    srl = StateRepresentation(obs_dim=[32, 24, 40], state_dim=5, act_dim=2, batch_size=256, learning_rate=5e-4, alpha=2,
                              beta=10, seed=3, continuelearning=False, usingros=False)

    # add all data to srl memory
    for d in data:
        srl.remember(reshape_observation(d[0]), d[1], d[2], d[3], d[4])

    epochs = 20
    loss_history = []

    # training the encoder for 20 epochs
    for epoch in range(epochs):
        loss, losses = srl.learn()
        print('Finished epoch {}/{}. The loss this epoch was: {}'.format(epoch + 1, epochs, loss))
        print('temp_coh_loss: {} caus_loss: {} prop_loss: {} repeat_loss: {}'.format(losses[1], losses[2], losses[3],
                                                                                losses[4]))

        loss_history.append(losses)

    # save the trained model
    srl.save_model()

    # compute all the states from the observations
    print('Finished training, now predicting all states')
    lg = Logger('./')
    gt = []
    for d in data:
        gt.append(d[0][40:45])
        lg.log("position_obs", d[0][40:45])
        lg.log("rewards", d[2])

    states, _ = srl.get_memory_states()

    # visualise the results
    print('Plotting the results')

    plotter.plot_all_srl(states=states, rews=srl.memory.get_all_samples().get('rews'), data_dict=lg.logDict,
                         loss_history=loss_history, trainingcycle=2, save=True, size=15000, name="continuous")


"""
adaptive split data into bins.
"""

import functools

import numpy as np

from .data import data_index, data_mask

try:
    import matplotlib.patches as mpathes
except ImportError:
    pass

class AdaptiveBound(object):
    """adaptive bound cut for data value"""

    def __init__(self, base_data, bins):
        if isinstance(bins, int):
            self._base_data = np.array([base_data])
            self.bins = [bins]
        elif isinstance(bins, list):
            self._base_data = np.array(base_data)
            self.bins = bins
        else:
            raise TypeError(
                "bins should be int or list of int, "
                + "insteading of {}".format(type(bins))
            )

    @functools.lru_cache()
    def get_bounds_data(self):
        """get split data bounds, and the data after splitting"""
        base_bound = AdaptiveBound.base_bound(self._base_data)
        bounds, datas = AdaptiveBound.loop_split_bound(
            self._base_data, self.bins
        )
        return bounds, datas

    def get_bounds(self):
        """get split data bounds"""
        bounds, _ = self.get_bounds_data()
        return bounds

    def get_bool_mask(self, data):
        """bool mask for splitting data"""
        bounds = self.get_bounds()
        idx_data = data[: bounds[0][0].shape[0]]
        ret = []
        for lb, rb in bounds:
            mask = np.logical_and(
                idx_data >= lb[..., np.newaxis], idx_data < rb[..., np.newaxis]
            )
            mask = np.all(mask, axis=0)
            ret.append(mask)
        return ret

    def split_full_data(self, data, base_index=None):
        """split structure data, (TODO because large IO,  the method is slow.)"""
        base_data = [[data_index(data, i) for i in base_index]]
        mask = self.get_bool_mask(base_data)
        ret = []
        for i in mask:
            ret.append(data_mask(data, i))
        return ret

    def split_data(self, data):
        """split data, the shape is same as base_data"""
        mask = self.get_bool_mask(data)
        ret = []
        for i in mask:
            ret.append(data[..., i])
        return ret

    @staticmethod
    def single_split_bound(data, n=2, base_bound=None):
        """split data in the order of data value

        >>> data = np.array([1.0, 2.0, 1.4, 3.1])
        >>> AdaptiveBound.single_split_bound(data)
        [(1.0, 1.7...), (1.7..., 3.1...)]

        """
        if base_bound is None:
            base_bound = np.min(data), np.max(data) + 1e-6
        num_lb = base_bound[0]
        bounds = []
        for j in range(1, n):
            num_rb = np.percentile(data, j / n * 100, axis=0) + 1e-6
            bounds.append((num_lb, num_rb))
            num_lb = num_rb
        bounds.append((num_lb, base_bound[1]))
        return bounds

    @staticmethod
    def multi_split_bound(datas, n, base_bound=None):
        """multi data for single_split_bound, so `n` is list of int

        >>> data = np.array([[1.0, 2.0, 1.4, 3.1], [2.0, 1.0, 3.0, 1.0]])
        >>> bound, _ = AdaptiveBound.multi_split_bound(data, [2, 1])
        >>> [(i[0][0]+1e-6, i[1][0]+1e-6) for i in bound]
        [(1.0..., 1.7...), (1.7..., 3.1...)]

        """
        datas = np.array(datas)
        if base_bound is None:
            base_bound = AdaptiveBound.base_bound(datas)
        bound_chain = [base_bound]
        data_chain = [datas]
        for idx, size in enumerate(n):
            new_bound_chain = []
            new_data_chain = []
            for bnd, data in zip(bound_chain, data_chain):
                idx_data = data[idx]
                bounds = AdaptiveBound.single_split_bound(
                    idx_data, size, base_bound=(bnd[0][idx], bnd[1][idx])
                )
                for i in bounds:
                    lb, rb = i
                    l_bnd, r_bnd = bnd
                    l_bnd = np.copy(l_bnd)
                    r_bnd = np.copy(r_bnd)
                    l_bnd[idx] = lb
                    r_bnd[idx] = rb
                    new_bound_chain.append((l_bnd, r_bnd))
                    mask = np.logical_and(idx_data >= lb, idx_data < rb)
                    new_data_chain.append(data[:, mask])
            bound_chain = new_bound_chain
            data_chain = new_data_chain
        return bound_chain, data_chain

    @staticmethod
    def loop_split_bound(datas, n, base_bound=None):
        """loop for multi_split_bound, so `n` is list of list of int"""
        datas = np.array(datas)
        if base_bound is None:
            base_bound = AdaptiveBound.base_bound(datas)
        bound_chain = [base_bound]
        data_chain = [datas]
        for idx, size in enumerate(n):
            new_bound_chain = []
            new_data_chain = []
            for bnd, data in zip(bound_chain, data_chain):
                bound, data_i = AdaptiveBound.multi_split_bound(
                    data, size, base_bound=bnd
                )
                new_bound_chain += bound
                new_data_chain += data_i
            bound_chain = new_bound_chain
            data_chain = new_data_chain
        return bound_chain, data_chain

    @staticmethod
    def base_bound(data):
        """base bound for the data"""
        lb = np.min(data, axis=-1) - 1e-6
        rb = np.max(data, axis=-1) + 1e-6
        return (lb, rb)

    def get_bound_patch(self, **kwargs):
        ret = []
        for i, bnd in enumerate(self.get_bounds()):
            min_x, min_y = bnd[0]
            max_x, max_y = bnd[1]
            rect = mpathes.Rectangle(
                (min_x, min_y), max_x - min_x, max_y - min_y, **kwargs
            )  # cmap(weights[i]/max_weight))
            ret.append(rect)
        return ret

    def plot_bound(self, ax, **kwargs):
        for i in self.get_bound_patch(**kwargs):
            ax.add_patch(i)

def cal_chi2(numbers, n_fp):
    weights = []
    # print(numbers)
    # chi21 = []
    for ndata, nmc in numbers:
        weight = (ndata - nmc) / np.sqrt(np.abs(ndata))
        weights.append(weight ** 2)
        # chi21.append(ndata * np.log(nmc))
    max_weight = np.max(weights)
    chi2 = np.sum(weights)
    print("bins: ", len(weights))
    print("number of free parameters: ", n_fp)
    ndf = len(weights) - 1 - n_fp
    print(
        "chi2/ndf: ", np.sum(weights), "/", ndf
    )  # ,"another", np.sum(chi21))
    return chi2, ndf


import psi4
import numpy as np

psi4.set_options({'scf_type': 'df',
                  'freeze_core': 'true'})

#different combinations of the isomers

#O/M = comparison1
#O/P = comparison2
#M/P = comparison3

#comparison1 O/M

comparison1 = psi4.geometry("""

0 1

O          1.72900        1.44890        1.30290
O          0.52150       -2.76860        0.50990
O          1.94280        1.60280       -0.95480
O          1.94610       -1.35340       -0.55510
C         -0.07460        0.73300       -0.00740
C         -0.28440       -0.64490       -0.05810
C         -1.16380        1.60440        0.00750
C         -1.58340       -1.15170       -0.09400
C         -2.46290        1.09750       -0.02840
C         -2.67270       -0.28050       -0.07920
C          1.27000        1.29310        0.03190
C          0.83250       -1.58060       -0.07500
H         -1.01340        2.68020        0.04680
H         -1.78090       -2.21920       -0.14740
H         -3.31100        1.77580       -0.01770
H         -3.68430       -0.67460       -0.11040
H          2.63430        1.82640        1.32040
H          1.27850       -3.39240        0.49470

--

0 1
O         -3.54720        0.18420        0.00030
O          3.54710        0.18450       -0.00070
O         -2.58510       -1.87520       -0.00050
O          2.58550       -1.87510        0.00060
C         -1.20810        0.08130        0.00030
C          1.20790        0.08110       -0.00010
C         -0.00020       -0.61630        0.00010
C         -1.20790        1.47630        0.00020
C          1.20790        1.47610       -0.00020
C          0.00000        2.17360        0.00010
C         -2.46790       -0.64530       -0.00040
C          2.46800       -0.64530        0.00040
H         -0.00020       -1.70390        0.00010
H         -2.12470        2.05900        0.00020
H          2.12470        2.05880       -0.00020
H          0.00010        3.25980        0.00000
H         -4.39140       -0.31530        0.00040
H          4.39150       -0.31500       -0.00100
""")

psi4.energy('sapt0/jun-cc-pvdz', molecule=comparison1)

one_disp = psi4.variable('SSAPT0 DISP ENERGY')
one_elst = psi4.variable('SSAPT0 ELST ENERGY')
one_exch = psi4.variable('SSAPT0 EXCH ENERGY')
one_ind = psi4.variable('SSAPT0 IND ENERGY')
one_tot =psi4.variable('SSAPT0 TOTAL ENERGY')

#comparison2 O/P

comparison2 = psi4.geometry("""

0 1

O          1.72900        1.44890        1.30290
O          0.52150       -2.76860        0.50990
O          1.94280        1.60280       -0.95480
O          1.94610       -1.35340       -0.55510
C         -0.07460        0.73300       -0.00740
C         -0.28440       -0.64490       -0.05810
C         -1.16380        1.60440        0.00750
C         -1.58340       -1.15170       -0.09400
C         -2.46290        1.09750       -0.02840
C         -2.67270       -0.28050       -0.07920
C          1.27000        1.29310        0.03190
C          0.83250       -1.58060       -0.07500
H         -1.01340        2.68020        0.04680
H         -1.78090       -2.21920       -0.14740
H         -3.31100        1.77580       -0.01770
H         -3.68430       -0.67460       -0.11040
H          2.63430        1.82640        1.32040
H          1.27850       -3.39240        0.49470

--

0 1
O          3.39660        1.18370       -0.00030
O         -3.39650       -1.18360       -0.00070
O          3.54230       -1.08460        0.00020
O         -3.54230        1.08470        0.00000
C          1.39460       -0.03050       -0.00050
C         -1.39450        0.03050        0.00050
C          0.72350        1.19240       -0.00030
C         -0.67100        1.22300        0.00030
C          0.67090       -1.22310       -0.00020
C         -0.72370       -1.19250        0.00030
C          2.84840       -0.06240        0.00030
C         -2.84820        0.06250        0.00030
H          1.24760        2.14410       -0.00030
H         -1.17530        2.18580        0.00040
H          1.17520       -2.18590       -0.00010
H         -1.24770       -2.14420        0.00050
H          4.37700        1.15110       -0.00040
H         -4.37700       -1.15080       -0.00130
""")

psi4.energy('sapt0/jun-cc-pvdz', molecule=comparison2)

two_disp = psi4.variable('SSAPT0 DISP ENERGY')
two_elst = psi4.variable('SSAPT0 ELST ENERGY')
two_exch = psi4.variable('SSAPT0 EXCH ENERGY')
two_ind = psi4.variable('SSAPT0 IND ENERGY')
two_tot =psi4.variable('SSAPT0 TOTAL ENERGY')

#comparison3 M/P

comparison3 = psi4.geometry("""

0 1
O         -3.54720        0.18420        0.00030
O          3.54710        0.18450       -0.00070
O         -2.58510       -1.87520       -0.00050
O          2.58550       -1.87510        0.00060
C         -1.20810        0.08130        0.00030
C          1.20790        0.08110       -0.00010
C         -0.00020       -0.61630        0.00010
C         -1.20790        1.47630        0.00020
C          1.20790        1.47610       -0.00020
C          0.00000        2.17360        0.00010
C         -2.46790       -0.64530       -0.00040
C          2.46800       -0.64530        0.00040
H         -0.00020       -1.70390        0.00010
H         -2.12470        2.05900        0.00020
H          2.12470        2.05880       -0.00020
H          0.00010        3.25980        0.00000
H         -4.39140       -0.31530        0.00040
H          4.39150       -0.31500       -0.00100

--

0 1
O          3.39660        1.18370       -0.00030
O         -3.39650       -1.18360       -0.00070
O          3.54230       -1.08460        0.00020
O         -3.54230        1.08470        0.00000
C          1.39460       -0.03050       -0.00050
C         -1.39450        0.03050        0.00050
C          0.72350        1.19240       -0.00030
C         -0.67100        1.22300        0.00030
C          0.67090       -1.22310       -0.00020
C         -0.72370       -1.19250        0.00030
C          2.84840       -0.06240        0.00030
C         -2.84820        0.06250        0.00030
H          1.24760        2.14410       -0.00030
H         -1.17530        2.18580        0.00040
H          1.17520       -2.18590       -0.00010
H         -1.24770       -2.14420        0.00050
H          4.37700        1.15110       -0.00040
H         -4.37700       -1.15080       -0.00130
""")

psi4.energy('sapt0/jun-cc-pvdz', molecule=comparison3)

three_disp = psi4.variable('SSAPT0 DISP ENERGY')
three_elst = psi4.variable('SSAPT0 ELST ENERGY')
three_exch = psi4.variable('SSAPT0 EXCH ENERGY')
three_ind = psi4.variable('SSAPT0 IND ENERGY')
three_tot =psi4.variable('SSAPT0 TOTAL ENERGY')

#comparing the interaction energies for each dimer

labels = np.array(['disp','elst','exch','ind','tot'])

dispersion = np.array([one_disp,two_disp,three_disp])

elst = np.array([one_elst,two_elst,three_elst])

exch = np.array([one_exch,two_exch,three_exch])

ind = np.array([one_ind,two_ind,three_ind])

ind = np.array([one_ind,two_ind,three_ind])

tot = np.array([one_tot,two_tot,three_tot])

#plotting these comparisons

plt.plot(comparison,tot,'-ob')

plt.plot(comparison,ind,'-xr')

plt.plot(comparison,exch,'-py')

plt.plot(comparison,elst,'-dk')

plt.plot(comparison,dispersion,'-m<')

plt.show()