# Lint as: python3
# Copyright 2020 Google 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
#
#     https://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.
"""Join a sequence of Ito Processes with specified correlations."""

import tensorflow.compat.v2 as tf

from tf_quant_finance.models import euler_sampling
from tf_quant_finance.models import generic_ito_process
from tf_quant_finance.models import ito_process
from tf_quant_finance.models import utils


class JoinedItoProcess(generic_ito_process.GenericItoProcess):
  """Join of Ito Processes with specified time dependent correlations.

  For a sequence of Ito processes `I_1, .., I_n` of dimensions `d_1,.., d_n`,
  the class initializes a process `I` of dimension `d_1 + .. + d_n` with
  marginal proceses `I_i` and a correlation function `Corr(t)`. That is, let the
  Ito Process `I_i` describe an SDE

  ```None
  dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dW^i_t
  ```

  where `a_i(t, x)` is a function taking values in `R^{d_i}`, `b_i(t, X_t)` is a
  function taking values in `d_i x d_i` matrices, `W_i` is a `d_i`-dimensional
  Brownian motion.

  Then `I` describes an SDE for the joint process `(X_1,..., X_n)` of dimension
  `d:= d_1 + ... + d_n`

  ```None
  dX^i = a_i(t, X^i_t) dt + b_i(t, X^i_t) dB^i_t,
  ```

  where `(B_1, ..., B_n) = chol(t) * (W_1, ..., W_n)` for a Cholesky
  decomposition `chol(t)` of the correlation matrix `Corr(t)` at time `t`.
  Here `(W_1, ..., W_n)` is `d`-dimensional vector and `Corr(t)` is a `d x d`
  correlation matrix.

  `Corr(t)` is represented as a block-diagonal formed of a list of matrices
  `[m_1(t), m_2(t), ..., m_k(t)]` with `sum(rank(m_i)) = d`.

  # Example. # Black-scholes and Heston model join.
  ```python
  import numpy as np
  import tensorflow as tf
  import tf_quant_finance as tff

  dtype = tf.float64
  # Define Black scholes model with zero rate and volatility `0.1`
  sigma = 0.1
  def drift_fn(t , x):
    return -sigma**2 / 2
  def vol_fn(t , x):
    return sigma * tf.ones([1, 1], dtype=x.dtype)
  black_scholes_process = tff.models.GenericItoProcess(
      dim=1, drift_fn=drift_fn, volatility_fn=vol_fn, dtype=dtype)
  # Define Heston Model
  epsilon = tff.math.piecewise.PiecewiseConstantFunc(
      jump_locations=[0.5], values=[0.01, 0.02], dtype=np.float64)
  heston_process = tff.models.heston_model.HestonModel(
      kappa=1.0, theta=0.04, epsilon=epsilon, rho=0.2, dtype=dtype)
  # Define join process where correlation between `black_scholes_process` and
  # log-asset price of the `heston_process` is 0.5.
  corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]
  # `corr_structure` corresponds to a `3 x 3` correlation matrix. Here Brownian
  # motion of `black_scholes_process` is correlated only with the 1st dimension
  # of `heston_process` but not with the second one.
  join_process = JoinedItoProcess(
      processes=[black_scholes_process, heston_process],
      corr_structure=corr_structure)
  # Sample 100,000 sample paths at times [0.1, 1.0] from the join process using
  # Sobol random sequence
  times = [0.1, 1.0]
  # Wrap sample_paths method with a tf.function
  sample_paths_fn = tf.function(process.sample_paths)
  samples = sample_paths_fn(
      times=times, time_step=0.01, num_samples=100000,
      initial_state=np.array([0.0, 1.0, 0.04]),
      random_type=random.RandomType.SOBOL)
  # Estimated correlations.
  np.corrcoef(samples[:, -1, :], rowvar=False)
  # Expected result:
  # [[1.        , 0.49567078, 0.08128067],
  #  [0.49567078, 1.        , 0.16580689],
  #  [0.08128067, 0.16580689, 1.        ]]
  ```
  """

  def __init__(self, processes, corr_structure, dtype=None, name=None):
    """Initializes a JoinedItoProcess.

    Takes a list of `processes` which are instances of `tff.models.ItoProcess`
    and a list `corr_structure` of correlation matrices and creates an Ito
    process that joins `processes` using the correlation structure.
    `corr_structure` describes block-diagonal structure of correlations for
    the Brownian motions in `processes`. For example, if the dimension of the
    JoinedItoProcess is `3` and
    `corr_structure = [[[1.0, 0.5], [0.5, 1.0]], [1.0]]`, then the introduced
    correlation is
    `Corr(t) = [[1.0, 0.5, 0.0], [0.5, 1.0, 0.0], [0.0, 0.0, 1.0]]`,
    where `Corr(t)` is the same as in the `JoinedItoProcess` docstring.

    Args:
      processes: A sequence of instances of `tff.models.ItoProcess`. All
        processes should have the same `dtype.`
      corr_structure: A list of correlation matrices. Each correlation matrix
        is either a `Tensor` of the same `dtype` as the `processes` and
        square shape  (i.e., `[d_i, d_i]` for some `d_i`) or a callable. The
        callables should accept a scalar (stands for time `t`) and return a
        square `Tensor`. The total dimension
        `sum([m.shape[-1] for m in corr_structure]` of correlation
        structure should be the same as the dimension of the `JoinedItoProcess`
        `sum([p.dim() for p in processes])`.
      dtype: The default `dtype` of the `processes`.
        Default value: None which means that default dtypes inferred by
          TensorFlow are used.
      name: Python string. The name scope under which ops created by the methods
        of this class are nested.
        Default value: `None` which maps to the default name
          `join_ito_process`.

    Raises:
      ValueError:
        (a) If any of the `processes` is not an `ItoProcess`.
        (b) If `processes` do not have the same `dtype`.
    """
    self._name = name or "join_ito_process"
    with tf.name_scope(self._name):
      self._processes = []  # Input processes
      dim = 0  # Dimension of the process
      for process in processes:
        if not isinstance(process, ito_process.ItoProcess):
          raise ValueError(
              "All input process of JoinedItoProcess must be instances "
              "of the ItoProcess class.")
        self._processes.append(process)
        d = process.dim()
        dim += d  # Total dimension of the process
        if dtype is None:
          dtype = process.dtype()
        elif dtype != process.dtype():
          raise ValueError("All processes should have the same `dtype`")
      self._corr_structure = [
          corr if callable(corr) else tf.convert_to_tensor(
              corr, dtype=dtype, name="corr")
          for corr in corr_structure]
      self._dim = dim

      def _drift_fn(t, x):
        """Drift function of the JoinedItoProcess."""
        drifts = []
        i1 = 0
        i2 = 0
        for p in self._processes:
          dim = p.dim()
          i2 += dim
          position = x[..., i1:i2]