# Copyright 2022 Battelle Energy Alliance, LLC
# ALL RIGHTS RESERVED
"""
Created on July 30 2020
@author: Congjian Wang
"""

#External Modules------------------------------------------------------------------------------------
import abc
import numpy as np
import logging
#External Modules End--------------------------------------------------------------------------------

#Internal Modules------------------------------------------------------------------------------------
from ravenframework.utils import mathUtils as utils
from ravenframework.utils import InputData, InputTypes
#Internal Modules End--------------------------------------------------------------------------------

#TODO: Use logging for errors and warnings
logging.basicConfig(format='%(asctime)s %(name)-20s %(levelname)-8s %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.DEBUG)
# To enable the logging to both file and console, the logger for the main should be the root,
# otherwise, a function to add the file handler and stream handler need to be created and called by each module.
# logger = logging.getLogger(__name__)
logger = logging.getLogger()
# # create file handler which logs debug messages
fh = logging.FileHandler(filename='log.log', mode='w')
fh.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s %(name)-20s %(levelname)-8s %(message)s')
fh.setFormatter(formatter)
# add the handlers to the logger
logger.addHandler(fh)


class LikelihoodBase:
  """
    Base class for Likelihood models
  """

  @classmethod
  def getInputSpecification(cls):
    """
      Collects input specifications for this class.

      :returns:
        InputData: RAVEN InputData specs
    """
    inputSpecs = InputData.parameterInputFactory('LikelihoodModel')
    inputSpecs.addParam('type', param_type=InputTypes.StringType,
        descr='The likelihood model object identifier')
    return inputSpecs

  def __init__(self):
    """
      Constructuor

      :returns: None
    """
    self.type = self.__class__.__name__
    self.name = self.__class__.__name__
    # dictionary: keys all required input parameters, and values either user provided values or
    # variables determined by raven
    self._variableDict = {}
    # instance of Likelihood model
    self._model = None
    # class of Likelihood model
    self._modelClass = None
    # outputs
    self._outputs = {}

  def handleInput(self, xmlNode):
    """
      Function to read the portion of the xml input that belongs to this specialized class
      and initialize some stuff based on the inputs got

      :param xmlNode: xml.etree.ElementTree.Element, Xml element node

      :returns: None
    """
    paramInput = self.getInputSpecification()()
    paramInput.parseNode(xmlNode)
    self._localHandleInput(paramInput)

  @abc.abstractmethod
  def _localHandleInput(self, paramInput):
    """
      Function to process the parsed xml input

      :param paramInput: InputData.ParameterInput, the parsed xml input

      :returns: None
    """

  def initialize(self, inputDict):
    """
      Method to initialize

      :param inputDict: dict, dictionary of inputs

      :returns: None
    """
    needDict = self.getParams()
    needDict = self.loadVariables(needDict, inputDict)
    self.setParams(needDict)
    self._checkInputParams(needDict)

  def _checkInputParams(self, needDict):
    """
      Method to check input parameters

      :param needDict: dict, dictionary of required parameters

      :returns: None
    """
    pass

  def setVariable(self, value):
    """
      Set value if a float/int/list is provided in the node text, othersise treat the provided value as RAVEN variable

      :param value: str or float or list, the value of given variable

      returns:
        str or numpy.array: the recasted value
    """
    ret = None
    # multi-entry or single-entry?
    if len(value) == 1:
      if not utils.isAFloatOrInt(value[0]):
        ret = value[0]
      else:
        ret = np.atleast_1d(value)
    else:
      # should be floats; InputData assures the entries are the same type already
      if not utils.isAFloatOrInt(value[0]):
        logger.warning('Multiple non-number entries are found, but require either a single variable \
                       name or multiple float entries: {}'.format(value))
        ret = value
      else:
        ret = np.asarray(value)
    return ret

  def loadVariables(self, need, inputDict):
    """
      Load the values of variables that is generated by RAVEN

      :param need: dict, the dict of parameters
      :param inputDict: dict, the dict of parameters that is provided from other sources

      returns:
        dict: the dict of parameters updated with variables
    """
    # load variable values from variables as needed
    for key, val in need.items():
      if utils.isAString(val):
        value = inputDict.get(val, None)
        if value is None:
          raise KeyError('Looking for variable "{}" to fill "{}" but not found among variables!'.format(val, key))
        need[key] = np.atleast_1d(value)
      # already converted list of floats into np.array, only need to convert list of strings
      elif type(val) == list:
        augmentedValue = []
        for v in val:
          value = inputDict.get(v, None)
          if value is None:
            raise KeyError('Looking for variable "{}" to fill "{}" but not found among variables!'.format(v, key))
          augmentedValue.extend(value)
        need[key] = np.atleast_1d(augmentedValue)
    return need

  def getParams(self):
    """
      Get the parameters

      returns:
        dict: dictionary of variables
    """
    return self._variableDict

  def setParams(self, paramDict):
    """
      Set the parameters from a given dictionary.

      :param paramDict: dict, settings

      :returns: None
    """
    for key, val in paramDict.items():
      if key in self.__dict__.keys():
        setattr(self, key, val)
      else:
        logger.warning('Variable "{}" is not defined in class "{}"!'.format(key, self.name))

  def getOutputs(self):
    """
      get calculated cdf value

      returns:
        dict: dictionary of model outputs
    """
    return self._outputs

  @abc.abstractmethod
  def _logLikelihoodFunction(self):
    """
      Function to calculate  log probability

      returns:
        dict: return log likelihood outputs
    """

  def run(self):
    """
      Method to calculate Likelihood related quantities

      :returns: None
    """
    self._outputs = self._logLikelihoodFunction()
