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

Contains a utility base class for accessing commonly-used TSA functions.
"""
import numpy as np

from utils import xmlUtils, InputData, InputTypes

from .Factory import factory

class TSAUser:
  """
    Add-on class for inheriting commonly-used TSA algorithms, such as reading in algorithms
  """
  @classmethod
  def addTSASpecs(cls, spec, subset=None):
    """
      Make input specs for TSA algorithm inclusion.
      @ In, spec, InputData.parameterInput, input specs to which TSA algos should be added
      @ In, subset, str, optional, one of (None, 'generate', 'characterize'), subset of allowable algorithms
      @ Out, None
    """
    # NOTE we use an assertion here, since only developer actions should be picking subType
    assert subset in (None, 'characterize', 'generate'), f'Invalid subset for TSAUser.addTSASpecs: "{subset}"'
    # need pivot parameter ID, but this might be superceded by a parent class.
    spec.addSub(InputData.parameterInputFactory('pivotParameter', contentType=InputTypes.StringType,
        descr=r"""If a time-dependent ROM is requested, please specifies the pivot
        variable (e.g. time, etc) used in the input HistorySet.""", default='time'))
    for typ in factory.knownTypes():
      c = factory.returnClass(typ)
      if subset == 'characterize' and not c.canCharacterize():
        continue
      elif subset == 'generate' and not c.canGenerate():
        continue
      spec.addSub(c.getInputSpecification())
    return spec

  def __init__(self):
    """
      Constructor.
      @ In, None
      @ Out, None
    """
    # delay import of factory to assure linear imports
    self._tsaAlgoSettings = {}       # initialization settings for each algorithm
    self._tsaTrainedParams = {}      # holds results of training each algorithm
    self._tsaAlgorithms = []         # list and order for TSA algorithms to use
    self.pivotParameterID = None     # string name for time-like pivot parameter # TODO base class?
    self.pivotParameterValues = None # values for the time-like pivot parameter  # TODO base class?
    self._paramNames = None          # cached list of parameter names
    self._paramRealization = None    # cached dict of param variables mapped to values
    self._tsaTargets = None          # cached list of targets
    self.target = None

  def readTSAInput(self, spec):
    """
      Read in TSA algorithms
      @ In, spec, InputData.parameterInput, input specs filled with user entries
      @ Out, None
    """
    if self.pivotParameterID is None: # might be handled by parent
      self.pivotParameterID = spec.findFirst('pivotParameter').value
    for sub in spec.subparts:
      if sub.name in factory.knownTypes():
        algo = factory.returnInstance(sub.name)
        self._tsaAlgoSettings[algo] = algo.handleInput(sub)
        self._tsaAlgorithms.append(algo)
        foundTSAType = True
    if foundTSAType is False:
      options = ', '.join(factory.knownTypes())
      # NOTE this assumes that every TSAUser is also an InputUser!
      raise IOError(f'TSA: No known TSA type found in input. Available options are: {options}')
    if self.target is None:
      # set up all the expected targets from all the TSAs
      self.target = [self.pivotParameterID] + list(self.getTargets())
    elif self.pivotParameterID not in self.target:
      # NOTE this assumes that every TSAUser is also an InputUser!
      raise IOError('TSA: The pivotParameter must be included in the target space.')

  def _tsaReset(self):
    """
      Resets trained and cached params
      @ In, None
      @ Out, None
    """
    self._tsaTrainedParams = {}      # holds results of training each algorithm
    self._paramNames = None          # cached list of parameter names
    self._paramRealization = None    # cached dict of param variables mapped to values
    self._tsaTargets = None          # cached list of targets

  def getTargets(self):
    """
      Provide ordered target set for the set of algorithms used by this entity
      @ In, None
      @ Out, targets, set, set of targets used among all algorithms
    """
    if self._tsaTargets is None:
      targets = set()
      # if we've trained params, we can use that
      if self._tsaTrainedParams:
        for algo, params in self._tsaTrainedParams.items():
          targets.update(params.keys())
      # otherwise, we use targets from settings
      else:
        for algo, settings in self._tsaAlgoSettings.items():
          targets.update(settings['target'])
      self._tsaTargets = targets
    return self._tsaTargets

  def getCharacterizingVariableNames(self):
    """
      Provide expected training parameters as variable names list
      Note this works even if training has not been performed yet
      @ In, None
      @ Out, names, list, list of parameter names
    """
    if self._paramNames is None:
      if self._paramRealization is not None:
        # if we trained already?,then we can return keys
        self._paramNames = list(self._paramRealization.keys())
      else:
        # otherwise we build the names predictively
        names = []
        for algo in self._tsaAlgorithms:
          names.extend(algo.getParamNames(self._tsaAlgoSettings[algo]))
        self._paramNames = names
    return self._paramNames

  def getParamsAsVars(self):
    """
      Provide training parameters as variable names mapped to values
      Note this is only useful AFTER characterization has been performed/trained
      @ In, None
      @ Out, params, dict, map of {algo_param: value}
    """
    if self._paramRealization is None:
      rlz = {}
      for algo, params in self._tsaTrainedParams.items():
        new = algo.getParamsAsVars(params)
        rlz.update(new)
      self._paramRealization = rlz
    return self._paramRealization

  def trainTSASequential(self, targetVals):
    """
      Train TSA algorithms using a sequential removal-and-residual approach.
      @ In, targetVals, array, shape = [n_timeStep, n_dimensions], array of time series data
        NOTE: this should be a single history/realization, not an array of realizations
      @ Out, None
    """
    pivotName = self.pivotParameterID
    # NOTE assumption: self.target exists!
    pivotIndex = self.target.index(pivotName)
    # NOTE assumption: only one training signal
    pivots = targetVals[0, :, pivotIndex]
    self.pivotParameterValues = pivots[:] # TODO any way to avoid storing these?
    residual = targetVals[:, :, :] # deep-ish copy, so we don't mod originals
    numAlgo = len(self._tsaAlgorithms)
    for a, algo in enumerate(self._tsaAlgorithms):
      settings = self._tsaAlgoSettings[algo]
      targets = settings['target']
      indices = tuple(self.target.index(t) for t in targets)
      signal = residual[0, :, indices].T # using tuple "indices" transposes, so transpose back
      params = algo.characterize(signal, pivots, targets, settings)
      # store characteristics
      self._tsaTrainedParams[algo] = params
      # obtain residual; the part of the signal not characterized by this algo
      # workaround: skip the last one, since it's often the ARMA and the residual isn't known for
      #             the ARMA
      if a < numAlgo - 1:
        algoResidual = algo.getResidual(signal, params, pivots, settings)
        residual[0, :, indices] = algoResidual.T # transpose, again because of indices
      # TODO meta store signal, residual?

  def evaluateTSASequential(self):
    """
      Evaluate TSA algorithms using a sequential linear superposition approach
      @ In, None
      @ Out, rlz, dict, realization dictionary of values for each target
    """
    pivots = self.pivotParameterValues
    # the algorithms' targets need to be consistently indexed, but there's no
    # reason to keep including the pivot values every time, so set up an indexing
    # that ignores the pivotParameter on which to index the results variables
    noPivotTargets = [x for x in self.target if x != self.pivotParameterID]
    result = np.zeros((self.pivotParameterValues.size, len(noPivotTargets)))
    for algo in self._tsaAlgorithms[::-1]:
      settings = self._tsaAlgoSettings[algo]
      targets = settings['target']
      indices = tuple(noPivotTargets.index(t) for t in targets)
      params = self._tsaTrainedParams[algo]
      if not algo.canGenerate():
        self.raiseAnError(IOError, "This TSA algorithm cannot generate synthetic histories.")
      signal = algo.generate(params, pivots, settings)
      result[:, indices] += signal
    # RAVEN realization construction
    rlz = dict((target, result[:, t]) for t, target in enumerate(noPivotTargets))
    rlz[self.pivotParameterID] = self.pivotParameterValues
    return rlz

  def writeTSAtoXML(self, xml):
    """
      Write properties of TSA algorithms to XML
      @ In, xml, xmlUtils.StaticXmlElement, entity to write to
      @ Out, None
    """
    root = xml.getRoot()
    for algo in self._tsaAlgorithms:
      algoNode = xmlUtils.newNode(algo.name)
      algo.writeXML(algoNode, self._tsaTrainedParams[algo])
      root.append(algoNode)

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

Contains a utility base class for accessing commonly-used TSA functions.
"""
import numpy as np

from utils import xmlUtils, InputData, InputTypes

from .Factory import factory

class TSAUser:
  """
    Add-on class for inheriting commonly-used TSA algorithms, such as reading in algorithms
  """
  @classmethod
  def addTSASpecs(cls, spec, subset=None):
    """
      Make input specs for TSA algorithm inclusion.
      @ In, spec, InputData.parameterInput, input specs to which TSA algos should be added
      @ In, subset, str, optional, one of (None, 'generate', 'characterize'), subset of allowable algorithms
      @ Out, None
    """
    # NOTE we use an assertion here, since only developer actions should be picking subType
    assert subset in (None, 'characterize', 'generate'), f'Invalid subset for TSAUser.addTSASpecs: "{subset}"'
    # need pivot parameter ID, but this might be superceded by a parent class.
    spec.addSub(InputData.parameterInputFactory('pivotParameter', contentType=InputTypes.StringType,
        descr=r"""If a time-dependent ROM is requested, please specifies the pivot
        variable (e.g. time, etc) used in the input HistorySet.""", default='time'))
    for typ in factory.knownTypes():
      c = factory.returnClass(typ)
      if subset == 'characterize' and not c.canCharacterize():
        continue
      elif subset == 'generate' and not c.canGenerate():
        continue
      spec.addSub(c.getInputSpecification())
    return spec

  def __init__(self):
    """
      Constructor.
      @ In, None
      @ Out, None
    """
    # delay import of factory to assure linear imports
    self._tsaAlgoSettings = {}       # initialization settings for each algorithm
    self._tsaTrainedParams = {}      # holds results of training each algorithm
    self._tsaAlgorithms = []         # list and order for TSA algorithms to use
    self.pivotParameterID = None     # string name for time-like pivot parameter # TODO base class?
    self.pivotParameterValues = None # values for the time-like pivot parameter  # TODO base class?
    self._paramNames = None          # cached list of parameter names
    self._paramRealization = None    # cached dict of param variables mapped to values
    self._tsaTargets = None          # cached list of targets
    self.target = None

  def readTSAInput(self, spec):
    """
      Read in TSA algorithms
      @ In, spec, InputData.parameterInput, input specs filled with user entries
      @ Out, None
    """
    if self.pivotParameterID is None: # might be handled by parent
      self.pivotParameterID = spec.findFirst('pivotParameter').value
    for sub in spec.subparts:
      if sub.name in factory.knownTypes():
        algo = factory.returnInstance(sub.name)
        self._tsaAlgoSettings[algo] = algo.handleInput(sub)
        self._tsaAlgorithms.append(algo)
        foundTSAType = True
    if foundTSAType is False:
      options = ', '.join(factory.knownTypes())
      # NOTE this assumes that every TSAUser is also an InputUser!
      raise IOError(f'TSA: No known TSA type found in input. Available options are: {options}')
    if self.target is None:
      # set up all the expected targets from all the TSAs
      self.target = [self.pivotParameterID] + list(self.getTargets())
    elif self.pivotParameterID not in self.target:
      # NOTE this assumes that every TSAUser is also an InputUser!
      raise IOError('TSA: The pivotParameter must be included in the target space.')

  def _tsaReset(self):
    """
      Resets trained and cached params
      @ In, None
      @ Out, None
    """
    self._tsaTrainedParams = {}      # holds results of training each algorithm
    self._paramNames = None          # cached list of parameter names
    self._paramRealization = None    # cached dict of param variables mapped to values
    self._tsaTargets = None          # cached list of targets

  def getTargets(self):
    """
      Provide ordered target set for the set of algorithms used by this entity
      @ In, None
      @ Out, targets, set, set of targets used among all algorithms
    """
    if self._tsaTargets is None:
      targets = set()
      # if we've trained params, we can use that
      if self._tsaTrainedParams:
        for algo, params in self._tsaTrainedParams.items():
          targets.update(params.keys())
      # otherwise, we use targets from settings
      else:
        for algo, settings in self._tsaAlgoSettings.items():
          targets.update(settings['target'])
      self._tsaTargets = targets
    return self._tsaTargets

  def getCharacterizingVariableNames(self):
    """
      Provide expected training parameters as variable names list
      Note this works even if training has not been performed yet
      @ In, None
      @ Out, names, list, list of parameter names
    """
    if self._paramNames is None:
      if self._paramRealization is not None:
        # if we trained already?,then we can return keys
        self._paramNames = list(self._paramRealization.keys())
      else:
        # otherwise we build the names predictively
        names = []
        for algo in self._tsaAlgorithms:
          names.extend(algo.getParamNames(self._tsaAlgoSettings[algo]))
        self._paramNames = names
    return self._paramNames

  def getParamsAsVars(self):
    """
      Provide training parameters as variable names mapped to values
      Note this is only useful AFTER characterization has been performed/trained
      @ In, None
      @ Out, params, dict, map of {algo_param: value}
    """
    if self._paramRealization is None:
      rlz = {}
      for algo, params in self._tsaTrainedParams.items():
        new = algo.getParamsAsVars(params)
        rlz.update(new)
      self._paramRealization = rlz
    return self._paramRealization

  def trainTSASequential(self, targetVals):
    """
      Train TSA algorithms using a sequential removal-and-residual approach.
      @ In, targetVals, array, shape = [n_timeStep, n_dimensions], array of time series data
        NOTE: this should be a single history/realization, not an array of realizations
      @ Out, None
    """
    pivotName = self.pivotParameterID
    # NOTE assumption: self.target exists!
    pivotIndex = self.target.index(pivotName)
    # NOTE assumption: only one training signal
    pivots = targetVals[0, :, pivotIndex]
    self.pivotParameterValues = pivots[:] # TODO any way to avoid storing these?
    residual = targetVals[:, :, :] # deep-ish copy, so we don't mod originals
    numAlgo = len(self._tsaAlgorithms)
    for a, algo in enumerate(self._tsaAlgorithms):
      settings = self._tsaAlgoSettings[algo]
      targets = settings['target']
      indices = tuple(self.target.index(t) for t in targets)
      signal = residual[0, :, indices].T # using tuple "indices" transposes, so transpose back
      params = algo.characterize(signal, pivots, targets, settings)
      # store characteristics
      self._tsaTrainedParams[algo] = params
      # obtain residual; the part of the signal not characterized by this algo
      # workaround: skip the last one, since it's often the ARMA and the residual isn't known for
      #             the ARMA
      if a < numAlgo - 1:
        algoResidual = algo.getResidual(signal, params, pivots, settings)
        residual[0, :, indices] = algoResidual.T # transpose, again because of indices
      # TODO meta store signal, residual?

  def evaluateTSASequential(self):
    """
      Evaluate TSA algorithms using a sequential linear superposition approach
      @ In, None
      @ Out, rlz, dict, realization dictionary of values for each target
    """
    pivots = self.pivotParameterValues
    # the algorithms' targets need to be consistently indexed, but there's no
    # reason to keep including the pivot values every time, so set up an indexing
    # that ignores the pivotParameter on which to index the results variables
    noPivotTargets = [x for x in self.target if x != self.pivotParameterID]
    result = np.zeros((self.pivotParameterValues.size, len(noPivotTargets)))
    for algo in self._tsaAlgorithms[::-1]:
      settings = self._tsaAlgoSettings[algo]
      targets = settings['target']
      indices = tuple(noPivotTargets.index(t) for t in targets)
      params = self._tsaTrainedParams[algo]
      if not algo.canGenerate():
        self.raiseAnError(IOError, "This TSA algorithm cannot generate synthetic histories.")
      signal = algo.generate(params, pivots, settings)
      result[:, indices] += signal
    # RAVEN realization construction
    rlz = dict((target, result[:, t]) for t, target in enumerate(noPivotTargets))
    rlz[self.pivotParameterID] = self.pivotParameterValues
    return rlz

  def writeTSAtoXML(self, xml):
    """
      Write properties of TSA algorithms to XML
      @ In, xml, xmlUtils.StaticXmlElement, entity to write to
      @ Out, None
    """
    root = xml.getRoot()
    for algo in self._tsaAlgorithms:
      algoNode = xmlUtils.newNode(algo.name)
      algo.writeXML(algoNode, self._tsaTrainedParams[algo])
      root.append(algoNode)

# Copyright 2017 Battelle Energy Alliance, LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. Created on August 3, 2021 @author: talbpaul  Contains a utility base class for accessing commonly-used TSA functions. Add-on class for inheriting commonly-used TSA algorithms, such as reading in algorithms Make input specs for TSA algorithm inclusion.       @ In, spec, InputData.parameterInput, input specs to which TSA algos should be added       @ In, subset, str, optional, one of (None, 'generate', 'characterize'), subset of allowable algorithms       @ Out, None # NOTE we use an assertion here, since only developer actions should be picking subType # need pivot parameter ID, but this might be superceded by a parent class. If a time-dependent ROM is requested, please specifies the pivot         variable (e.g. time, etc) used in the input HistorySet. Constructor.       @ In, None       @ Out, None # delay import of factory to assure linear imports # initialization settings for each algorithm # holds results of training each algorithm # list and order for TSA algorithms to use # string name for time-like pivot parameter # TODO base class? # values for the time-like pivot parameter  # TODO base class? # cached list of parameter names # cached dict of param variables mapped to values # cached list of targets Read in TSA algorithms       @ In, spec, InputData.parameterInput, input specs filled with user entries       @ Out, None # might be handled by parent # NOTE this assumes that every TSAUser is also an InputUser! # set up all the expected targets from all the TSAs # NOTE this assumes that every TSAUser is also an InputUser! Resets trained and cached params       @ In, None       @ Out, None # holds results of training each algorithm # cached list of parameter names # cached dict of param variables mapped to values # cached list of targets Provide ordered target set for the set of algorithms used by this entity       @ In, None       @ Out, targets, set, set of targets used among all algorithms # if we've trained params, we can use that # otherwise, we use targets from settings Provide expected training parameters as variable names list       Note this works even if training has not been performed yet       @ In, None       @ Out, names, list, list of parameter names # if we trained already?,then we can return keys # otherwise we build the names predictively Provide training parameters as variable names mapped to values       Note this is only useful AFTER characterization has been performed/trained       @ In, None       @ Out, params, dict, map of {algo_param: value} Train TSA algorithms using a sequential removal-and-residual approach.       @ In, targetVals, array, shape = [n_timeStep, n_dimensions], array of time series data         NOTE: this should be a single history/realization, not an array of realizations       @ Out, None # NOTE assumption: self.target exists! # NOTE assumption: only one training signal # TODO any way to avoid storing these? # deep-ish copy, so we don't mod originals # using tuple "indices" transposes, so transpose back # store characteristics # obtain residual; the part of the signal not characterized by this algo # workaround: skip the last one, since it's often the ARMA and the residual isn't known for #             the ARMA # transpose, again because of indices # TODO meta store signal, residual? Evaluate TSA algorithms using a sequential linear superposition approach       @ In, None       @ Out, rlz, dict, realization dictionary of values for each target # the algorithms' targets need to be consistently indexed, but there's no # reason to keep including the pivot values every time, so set up an indexing # that ignores the pivotParameter on which to index the results variables # RAVEN realization construction Write properties of TSA algorithms to XML       @ In, xml, xmlUtils.StaticXmlElement, entity to write to       @ Out, None