"""
Generators for some classic graphs.

The typical graph generator is called as follows:

>>> G=nx.complete_graph(100)

returning the complete graph on n nodes labeled 0,..,99
as a simple graph. Except for empty_graph, all the generators
in this module return a Graph class (i.e. a simple, undirected graph).

"""
#    Copyright (C) 2004-2010 by
#    Aric Hagberg <hagberg@lanl.gov>
#    Dan Schult <dschult@colgate.edu>
#    Pieter Swart <swart@lanl.gov>
#    All rights reserved.
#    BSD license.
import itertools
from networkx.algorithms.bipartite.generators import complete_bipartite_graph
__author__ ="""Aric Hagberg (hagberg@lanl.gov)\nPieter Swart (swart@lanl.gov)"""

__all__ = [ 'balanced_tree',
            'barbell_graph',
            'complete_graph',
            'circular_ladder_graph',
            'cycle_graph',
            'dorogovtsev_goltsev_mendes_graph',
            'empty_graph',
            'full_rary_tree',
            'grid_graph',
            'grid_2d_graph',
            'hypercube_graph',
            'ladder_graph',
            'lollipop_graph',
            'null_graph',
            'path_graph',
            'star_graph',
            'trivial_graph',
            'wheel_graph']

#-------------------------------------------------------------------
#   Some Classic Graphs
#-------------------------------------------------------------------
import networkx as nx
from networkx.utils import is_list_of_ints, flatten

def _tree_edges(n,r):
    # helper function for trees
    # yields edges in rooted tree at 0 with n nodes and branching ratio r
    nodes=iter(range(n))
    parents=[next(nodes)] # stack of max length r
    while parents:
        source=parents.pop(0)
        for i in range(r):
            try:
                target=next(nodes)
                parents.append(target)
                yield source,target
            except StopIteration:
                break

def full_rary_tree(r, n, create_using=None):
    """Creates a full r-ary tree of n vertices.

    Sometimes called a k-ary, n-ary, or m-ary tree.  "... all non-leaf
    vertices have exactly r children and all levels are full except
    for some rightmost position of the bottom level (if a leaf at the
    bottom level is missing, then so are all of the leaves to its
    right." [1]_

    Parameters
    ----------
    r : int
        branching factor of the tree
    n : int
        Number of nodes in the tree
    create_using : NetworkX graph type, optional
        Use specified type to construct graph (default = networkx.Graph)

    Returns
    -------
    G : networkx Graph
        An r-ary tree with n nodes

    References
    ----------
    .. [1] An introduction to data structures and algorithms,
           James Andrew Storer,  Birkhauser Boston 2001, (page 225).
    """
    G=nx.empty_graph(n,create_using)
    G.add_edges_from(_tree_edges(n,r))
    return G

def balanced_tree(r, h, create_using=None):
    """Return the perfectly balanced r-tree of height h.

    Parameters
    ----------
    r : int
        Branching factor of the tree
    h : int
        Height of the tree
    create_using : NetworkX graph type, optional
        Use specified type to construct graph (default = networkx.Graph)

    Returns
    -------
    G : networkx Graph
        A tree with n nodes

    Notes
    -----
    This is the rooted tree where all leaves are at distance h from
    the root. The root has degree r and all other internal nodes have
    degree r+1.

    Node labels are the integers 0 (the root) up to  number_of_nodes - 1.

    Also refered to as a complete r-ary tree.
    """
    # number of nodes is n=1+r+..+r^h
    if r==1:
        n=2
    else:
        n = int((1-r**(h+1))/(1-r)) # sum of geometric series r!=1
    G=nx.empty_graph(n,create_using)
    G.add_edges_from(_tree_edges(n,r))
    return G

    return nx.full_rary_tree(r,n,create_using)

def barbell_graph(m1,m2,create_using=None):
    """Return the Barbell Graph: two complete graphs connected by a path.

    For m1 > 1 and m2 >= 0.

    Two identical complete graphs K_{m1} form the left and right bells,
    and are connected by a path P_{m2}.

    The 2*m1+m2  nodes are numbered
        0,...,m1-1 for the left barbell,
        m1,...,m1+m2-1 for the path,
        and m1+m2,...,2*m1+m2-1 for the right barbell.

    The 3 subgraphs are joined via the edges (m1-1,m1) and (m1+m2-1,m1+m2).
    If m2=0, this is merely two complete graphs joined together.

    This graph is an extremal example in David Aldous
    and Jim Fill's etext on Random Walks on Graphs.

    """
    if create_using is not None and create_using.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")
    if m1<2:
        raise nx.NetworkXError(\
              "Invalid graph description, m1 should be >=2")
    if m2<0:
        raise nx.NetworkXError(\
              "Invalid graph description, m2 should be >=0")

    # left barbell
    G=complete_graph(m1,create_using)
    G.name="barbell_graph(%d,%d)"%(m1,m2)

    # connecting path
    G.add_nodes_from([v for v in range(m1,m1+m2-1)])
    if m2>1:
        G.add_edges_from([(v,v+1) for v in range(m1,m1+m2-1)])
    # right barbell
    G.add_edges_from( (u,v) for u in range(m1+m2,2*m1+m2) for v in range(u+1,2*m1+m2))
    # connect it up
    G.add_edge(m1-1,m1)
    if m2>0:
        G.add_edge(m1+m2-1,m1+m2)
    return G

def complete_graph(n,create_using=None):
    """ Return the complete graph K_n with n nodes.

    Node labels are the integers 0 to n-1.
    """
    G=empty_graph(n,create_using)
    G.name="complete_graph(%d)"%(n)
    if n>1:
        if G.is_directed():
            edges=itertools.permutations(range(n),2)
        else:
            edges=itertools.combinations(range(n),2)
        G.add_edges_from(edges)
    return G

def circular_ladder_graph(n,create_using=None):
    """Return the circular ladder graph CL_n of length n.

    CL_n consists of two concentric n-cycles in which
    each of the n pairs of concentric nodes are joined by an edge.

    Node labels are the integers 0 to n-1

    """
    G=ladder_graph(n,create_using)
    G.name="circular_ladder_graph(%d)"%n
    G.add_edge(0,n-1)
    G.add_edge(n,2*n-1)
    return G

def cycle_graph(n,create_using=None):
    """Return the cycle graph C_n over n nodes.

    C_n is the n-path with two end-nodes connected.

    Node labels are the integers 0 to n-1
    If create_using is a DiGraph, the direction is in increasing order.

    """
    G=path_graph(n,create_using)
    G.name="cycle_graph(%d)"%n
    if n>1: G.add_edge(n-1,0)
    return G

def dorogovtsev_goltsev_mendes_graph(n,create_using=None):
    """Return the hierarchically constructed Dorogovtsev-Goltsev-Mendes graph.

    n is the generation.
    See: arXiv:/cond-mat/0112143 by Dorogovtsev, Goltsev and Mendes.

    """
    if create_using is not None:
        if create_using.is_directed():
            raise nx.NetworkXError("Directed Graph not supported")
        if create_using.is_multigraph():
            raise nx.NetworkXError("Multigraph not supported")
    G=empty_graph(0,create_using)
    G.name="Dorogovtsev-Goltsev-Mendes Graph"
    G.add_edge(0,1)
    if n==0:
        return G
    new_node = 2         # next node to be added
    for i in range(1,n+1): #iterate over number of generations.
        last_generation_edges = G.edges()
        number_of_edges_in_last_generation = len(last_generation_edges)
        for j in range(0,number_of_edges_in_last_generation):
            G.add_edge(new_node,last_generation_edges[j][0])
            G.add_edge(new_node,last_generation_edges[j][1])
            new_node += 1
    return G

def empty_graph(n=0,create_using=None):
    """Return the empty graph with n nodes and zero edges.

    Node labels are the integers 0 to n-1

    For example:
    >>> G=nx.empty_graph(10)
    >>> G.number_of_nodes()
    10
    >>> G.number_of_edges()
    0

    The variable create_using should point to a "graph"-like object that
    will be cleaned (nodes and edges will be removed) and refitted as
    an empty "graph" with n nodes with integer labels. This capability
    is useful for specifying the class-nature of the resulting empty
    "graph" (i.e. Graph, DiGraph, MyWeirdGraphClass, etc.).

    The variable create_using has two main uses:
    Firstly, the variable create_using can be used to create an
    empty digraph, network,etc.  For example,

    >>> n=10
    >>> G=nx.empty_graph(n,create_using=nx.DiGraph())

    will create an empty digraph on n nodes.

    Secondly, one can pass an existing graph (digraph, pseudograph,
    etc.) via create_using. For example, if G is an existing graph
    (resp. digraph, pseudograph, etc.), then empty_graph(n,create_using=G)
    will empty G (i.e. delete all nodes and edges using G.clear() in
    base) and then add n nodes and zero edges, and return the modified
    graph (resp. digraph, pseudograph, etc.).

    See also create_empty_copy(G).

    """
    if create_using is None:
        # default empty graph is a simple graph
        G=nx.Graph()
    else:
        G=create_using
        G.clear()

    G.add_nodes_from(range(n))
    G.name="empty_graph(%d)"%n
    return G

def grid_2d_graph(m,n,periodic=False,create_using=None):
    """ Return the 2d grid graph of mxn nodes,
        each connected to its nearest neighbors.
        Optional argument periodic=True will connect
        boundary nodes via periodic boundary conditions.
    """
    G=empty_graph(0,create_using)
    G.name="grid_2d_graph"
    rows=range(m)
    columns=range(n)
    G.add_nodes_from( (i,j) for i in rows for j in columns )
    G.add_edges_from( ((i,j),(i-1,j)) for i in rows for j in columns if i>0 )
    G.add_edges_from( ((i,j),(i,j-1)) for i in rows for j in columns if j>0 )
    if G.is_directed():
        G.add_edges_from( ((i,j),(i+1,j)) for i in rows for j in columns if i<m-1 )
        G.add_edges_from( ((i,j),(i,j+1)) for i in rows for j in columns if j<n-1 )
    if periodic:
        if n>2:
            G.add_edges_from( ((i,0),(i,n-1)) for i in rows )
            if G.is_directed():
                G.add_edges_from( ((i,n-1),(i,0)) for i in rows )
        if m>2:
            G.add_edges_from( ((0,j),(m-1,j)) for j in columns )
            if G.is_directed():
                G.add_edges_from( ((m-1,j),(0,j)) for j in columns )
        G.name="periodic_grid_2d_graph(%d,%d)"%(m,n)
    return G

def grid_graph(dim,periodic=False):
    """ Return the n-dimensional grid graph.

    The dimension is the length of the list 'dim' and the
    size in each dimension is the value of the list element.

    E.g. G=grid_graph(dim=[2,3]) produces a 2x3 grid graph.

    If periodic=True then join grid edges with periodic boundary conditions.

    """
    dlabel="%s"%dim
    if dim==[]:
        G=empty_graph(0)
        G.name="grid_graph(%s)"%dim
        return G
    if not is_list_of_ints(dim):
        raise nx.NetworkXError("dim is not a list of integers")
    if min(dim)<=0:
        raise nx.NetworkXError(\
              "dim is not a list of strictly positive integers")
    if periodic:
        func=cycle_graph
    else:
        func=path_graph

    dim=list(dim)
    current_dim=dim.pop()
    G=func(current_dim)
    while len(dim)>0:
        current_dim=dim.pop()
        # order matters: copy before it is cleared during the creation of Gnew
        Gold=G.copy()
        Gnew=func(current_dim)
        # explicit: create_using=None
        # This is so that we get a new graph of Gnew's class.
        G=nx.cartesian_product(Gnew,Gold)
    # graph G is done but has labels of the form (1,(2,(3,1)))
    # so relabel
    H=nx.relabel_nodes(G, flatten)
    H.name="grid_graph(%s)"%dlabel
    return H

def hypercube_graph(n):
    """Return the n-dimensional hypercube.

    Node labels are the integers 0 to 2**n - 1.

    """
    dim=n*[2]
    G=grid_graph(dim)
    G.name="hypercube_graph_(%d)"%n
    return G

def ladder_graph(n,create_using=None):
    """Return the Ladder graph of length n.

    This is two rows of n nodes, with
    each pair connected by a single edge.

    Node labels are the integers 0 to 2*n - 1.

    """
    if create_using is not None and create_using.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")
    G=empty_graph(2*n,create_using)
    G.name="ladder_graph_(%d)"%n
    G.add_edges_from([(v,v+1) for v in range(n-1)])
    G.add_edges_from([(v,v+1) for v in range(n,2*n-1)])
    G.add_edges_from([(v,v+n) for v in range(n)])
    return G

def lollipop_graph(m,n,create_using=None):
    """Return the Lollipop Graph; `K_m` connected to `P_n`.

    This is the Barbell Graph without the right barbell.

    For m>1 and n>=0, the complete graph K_m is connected to the
    path P_n.  The resulting m+n nodes are labelled 0,...,m-1 for the
    complete graph and m,...,m+n-1 for the path. The 2 subgraphs
    are joined via the edge (m-1,m).  If n=0, this is merely a complete
    graph.

    Node labels are the integers 0 to number_of_nodes - 1.

    (This graph is an extremal example in David Aldous and Jim
    Fill's etext on Random Walks on Graphs.)

    """
    if create_using is not None and create_using.is_directed():
        raise nx.NetworkXError("Directed Graph not supported")
    if m<2:
        raise nx.NetworkXError(\
              "Invalid graph description, m should be >=2")
    if n<0:
        raise nx.NetworkXError(\
              "Invalid graph description, n should be >=0")
    # the ball
    G=complete_graph(m,create_using)
    # the stick
    G.add_nodes_from([v for v in range(m,m+n)])
    if n>1:
        G.add_edges_from([(v,v+1) for v in range(m,m+n-1)])
    # connect ball to stick
    if m>0: G.add_edge(m-1,m)
    G.name="lollipop_graph(%d,%d)"%(m,n)
    return G

def null_graph(create_using=None):
    """Return the Null graph with no nodes or edges.

    See empty_graph for the use of create_using.

    """
    G=empty_graph(0,create_using)
    G.name="null_graph()"
    return G

def path_graph(n,create_using=None):
    """Return the Path graph P_n of n nodes linearly connected by n-1 edges.

    Node labels are the integers 0 to n - 1.
    If create_using is a DiGraph then the edges are directed in
    increasing order.

    """
    G=empty_graph(n,create_using)
    G.name="path_graph(%d)"%n
    G.add_edges_from([(v,v+1) for v in range(n-1)])
    return G

def star_graph(n,create_using=None):
    """ Return the Star graph with n+1 nodes: one center node, connected to n outer nodes.

   Node labels are the integers 0 to n.

    """
    G=complete_bipartite_graph(1,n,create_using)
    G.name="star_graph(%d)"%n
    return G

def trivial_graph(create_using=None):
    """ Return the Trivial graph with one node (with integer label 0) and no edges.

    """
    G=empty_graph(1,create_using)
    G.name="trivial_graph()"
    return G

def wheel_graph(n,create_using=None):
    """ Return the wheel graph: a single hub node connected to each node of the (n-1)-node cycle graph.

   Node labels are the integers 0 to n - 1.

    """
    if n == 0:
        return nx.empty_graph(n, create_using=create_using)
    G=star_graph(n-1,create_using)
    G.name="wheel_graph(%d)"%n
    G.add_edges_from([(v,v+1) for v in range(1,n-1)])
    if n>2:
        G.add_edge(1,n-1)
    return G


# -*- coding: utf-8 -*-

# Copyright (c) 2012, Sergio Callegari
# All rights reserved.

# This file is part of PyDSM.

# PyDSM is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# PyDSM is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with PyDSM.  If not, see <http://www.gnu.org/licenses/>.

"""
Compatibility functions for PyDSM (:mod:`pydsm.relab`)
======================================================

This module re-implements some Matlab interfaces that are useful
for PyDSM.

.. currentmodule:: pydsm.relab

Functions
---------

.. autosummary::
   :toctree: generated/

   eps       -- Floating point relative accuracy
   db        -- Converts a value to dB a la Matlab
   cplxpair  -- Sorts values in input list by complex pairs
   shiftdim  -- Shift dimensions a la Matlab
"""

from __future__ import division, print_function

import numpy as np

__all__ = ["eps", "db", "cplxpair", "shiftdim"]

def eps(x=1.):
    """Provide floating point relative accuracy

    This function tries to replicate the ``eps`` interface of Matlab.

    Parameters
    ----------
    x : float or array_like or numeric type

    Returns
    -------
    d : float or array_like
        If x is a float, the positive distance from abs(x) to the next larger
        in magnitude floating point number of the same precision as x.

        If x is an array, operation is elementwise.

        If x is a numeric type, operation is for 1.0 in that numeric type.

    Examples
    --------
    >>> eps()
    2.2204460492503131e-16
    >>> eps(1/2)
    1.1102230246251565e-16
    >>> eps(np.float32)
    1.1920929e-07
    """

    def _eps(xi):
        return np.finfo(xi).eps * np.abs(xi)
    if isinstance(x, (type, np.dtype)):
        return np.finfo(x).eps
    elif np.isscalar(x):
        return _eps(x)
    x = np.asarray(x)
    d = np.empty_like(x)
    d.flat = [_eps(xi) for xi in x.flat]
    return d

def db(x, signal_type='voltage', R=1):
    """
    Converts a value to dB a la Matlab

    This function tries to replicate the ``dB`` interface of Matlab.

    Parameters
    ----------
    x : real
        value to be converted. Should be positive and non null.
    signal_type : string, optional
        either 'voltage' or 'power'. Defaults to 'voltage'
    R : real, optional
        load resistance value. Defaults to 1.

    Returns
    -------
    y : real
        value in dB corresponding to x
        if signal_type is 'power' the result is 10*log10(x). Otherwise (if
        signal_type is 'voltage'), then power is measured over resistor R

    Notes
    -----
    The default R value assures that when signal_type is 'voltage' dB defaults
    to the classical 20*log10(x) computation.
    """
    if signal_type == 'power':
        return 10*np.log10(x)
    else:
        return 10*np.log10(np.abs(x)**2./R)

def shiftdim(x, n=None, nargout=2):
    """
    Shift dimensions a la Matlab

    When n is provided, shiftdim shifts the axes of x by n.
    If n is positive, it shifts the axes to the left, wrapping the
    leading axes with non unitary length to the end.
    When n is negative, it shifts the axes to the right, inserting n leading
    axes with unitary length.
    When n is not provided or None, it shifts the axes to the left, reducing
    the number of dimensions and removing all the leading axes with unitary
    length.

    Parameters
    ----------
    x : array like
        multi-dimensional array to operate upon
    n : int or None, optional
        amount to shift. Defaults to None, which means automatic computation
    nargout : int
        number of output values

    Returns
    -------
    y : ndarray
        the result of the axes shift operation
    n : int
        the actual shift

    Examples
    --------
    >>> from numpy.random import rand
    >>> a = rand(1, 1, 3, 1, 2)
    >>> b, n = shiftdim(a)
    >>> b.shape
    (3, 1, 2)
    >>> n
    2
    >>> c = shiftdim(b, -n, nargout=1)
    >>> np.alltrue(c == a)
    True
    >>> d = shiftdim(a, 3, nargout=1)
    >>> d.shape
    (1, 2, 1, 1, 3)

    >>> b, n = shiftdim([[[1]]])
    >>> b, n
    (array([[[1]]]), 0)
    """
    outsel = slice(nargout) if nargout > 1 else 0
    x = np.asarray(x)
    s = x.shape
    m = next((i for i, v in enumerate(s) if v > 1), 0)
    if n is None:
        n = m
    if n > 0:
        n = n % x.ndim
    if n > 0:
        if n <= m:
            x = x.reshape(s[n:])
        else:
            x = x.transpose(np.roll(range(x.ndim), -n))
    elif n < 0:
            x = x.reshape((1,)*(-n)+x.shape)
    return (x, n)[outsel]

def cplxpair(x, tol=None, dim=None):
    """
    Sorts values into complex pairs a la Matlab.

    The function takes a vector or multidimensional array of of complex
    conjugate pairs or real numbers and rearranges it so that the complex
    numbers are collected into matched pairs of complex conjugates. The pairs
    are ordered by increasing real part, with purely real elements placed
    after all the complex pairs.

    In the search for complex conjugate pairs a relative tolerance equal to
    ``tol`` is used for comparison purposes. The default tolerance is
    100 times the system floating point accuracy.

    If the input vector is a multidimensional array, the rearrangement is done
    working along the axis specifid by the parameter ``dim`` or along the
    first axis with non-unitary length if ``dim`` is not provided.

    Parameters
    ----------
    x : array_like of complex
        x is an array of complex values, with the assumption that it contains
        either real values or complex values in conjugate pairs.
    tol: real, optional
        relative tolerance for the recognition of pairs.
        Defaults to 100 times the system floating point accuracy for the
        specific number type.
    dim: integer, optional
        The axis to operate upon.

    Returns
    -------
    y : ndarray
        y is an array of complex values, with the same values in x, yet now
        sorted as complex pairs by increasing real part. Real elements in x
        are place after the complex pairs, sorted in increasing order.

    Raises
    ------
    ValueError
        'Complex numbers cannot be paired' if there are unpaired complex
        entries in x.

    Examples
    --------
    >>> a = np.exp(2j*np.pi*np.arange(0, 5)/5)
    >>> b1 = cplxpair(a)
    >>> b2 = np.asarray([-0.80901699-0.58778525j, -0.80901699+0.58778525j,
    ...                   0.30901699-0.95105652j,  0.30901699+0.95105652j,
    ...                   1.00000000+0.j])
    >>> np.allclose(b1, b2)
    True

    >>> cplxpair(1)
    array([1])

    >>> cplxpair([[5, 6, 4], [3, 2, 1]])
    array([[3, 2, 1],
           [5, 6, 4]])

    >>> cplxpair([[5, 6, 4], [3, 2, 1]], dim=1)
    array([[4, 5, 6],
           [1, 2, 3]])

    See also
    --------
    eps : the system floating point accuracy
    """

    def cplxpair_vec(x, tol):
        real_mask = np.abs(x.imag) <= tol*np.abs(x)
        x_real = np.sort(np.real(x[real_mask]))
        x_cplx = np.sort(x[np.logical_not(real_mask)])
        if x_cplx.size == 0:
            return x_real
        if (x_cplx.size % 2) != 0:
            raise ValueError('Complex numbers cannot be paired')
        if np.any(np.real(x_cplx[1::2])-np.real(x_cplx[0::2]) >
                  tol*np.abs(x_cplx[0::2])):
            raise ValueError('Complex numbers cannot be paired')
        start = 0
        while start < x_cplx.size:
            sim_len = next((i for i, v in enumerate(x_cplx[start+1:]) if
                           (np.abs(np.real(v)-np.real(x_cplx[start])) >
                            tol*np.abs(v))), x_cplx.size-start-1)+1
            if (sim_len % 2) != 0:
                sim_len -= 1
            # At this point, sim_len elements with identical real part
            # have been identified.
            sub_x = x_cplx[start:start+sim_len]
            srt = np.argsort(np.imag(sub_x))
            sub_x = sub_x[srt]
            if np.any(np.abs(np.imag(sub_x)+np.imag(sub_x[::-1])) >
                      tol*np.abs(sub_x)):
                raise ValueError('Complex numbers cannot be paired')
            # Output should contain "perfect" pairs. Hence, keep entries
            # with positive imaginary parts amd use conjugate for pair
            x_cplx[start:start+sim_len] = np.concatenate(
                (np.conj(sub_x[:sim_len//2-1:-1]),
                 sub_x[:sim_len//2-1:-1]))
            start += sim_len
        return np.concatenate((x_cplx, x_real))

    x = np.atleast_1d(x)
    if x.size == 0:
        return x
    if dim is None:
        dim = next((i for i, v in enumerate(x.shape) if v > 1), 0)
    if tol is None:
        try:
            tol = 100*eps(x.dtype)
        except:
            tol = 100*eps(np.float)
    return np.apply_along_axis(cplxpair_vec, dim, x, tol)


#!/usr/bin/env python
# coding: utf-8

# # Quant Analyst Intern Assessment
#
# ## Section #2: Multivariate Time Series Analysis
#
# ### In this section, you will build a Multivariate Time Series Model. Please use Jupyter Notebook for this exercise. You may answer all questions in MARKDOWN inside the Jupyter Notebook.

#
# ### a. Propose up to five factors that may affect the USD/CAD trading range and explain your reasons for making these assumptions. Also explain the type of data frequency you will use (for example, intraday 5 minutes/daily/weekly).
#
# *I choose the following four factors:
# WTI crude oil price, xauusd, U.S. retail sales year-on-year, U.S. crude oil exports to Canada.*
#
# *I think there are not many factors that affect USD/CAD. It is mainly considered from the oil price data and U.S. economic data.*
#
# *First of all, the Canadian dollar is a commodity currency. Canada’s economic orientation is also dominated by exports. It is one of the important oil exporting countries. Therefore, rising oil prices may put a lot of pressure on the Canadian dollar. At this time, its purchasing power is limited. It will be greatly reduced, and it will depreciate against major currencies. It should be noted that although Canada exports a lot of oil, it also imports a lot of oil from the United States. A considerable part of this is that Canada first exported crude oil to the United States, then import it back from the United States. So I chose the WTI crude oil price and the amount of crude oil exported from the U.S. to Canada.*
#
# *Secondly, US consumption data, QE, Taper, and fiscal policies will all have an impact on usd/cad. Gold is closely related to the US dollar, so I chose xauusd. Since consumption accounts for a large proportion of the US economy, so from the perspective of consumption data I chose the retail sales year-on-year, and similarly, I can also choose the consumer confidence index instead. Here is a detail, the poor employment data of the United States will drive the Canadian dollar to rise, because the United States is full of employment, and the production of shale oil also increases, the oil prices fall.*
#
# *The specific choice actually depends on the strategy in different periods. For example, the current international oil prices are rising, but Canadian export oil is selling well, causing usdcad to fall a lot. If the U.S. dollar taper, international oil prices are weak, OPEC increases production capacity, and the U.S. dollar appreciates, the Canadian dollar will weaken relatively, and the logic will become dependent on US consumption data and manufacturing data.*
#
# *In the actual trading strategy, the higher the data frequency, the better. The FX leverage ratio is very high, so it is trading back and forth in the min range, and no one can bear the fluctuation of the next day.*
#
# *And I think the best construction of this model should be based on the long-term trend of the U.S. dollar index and the long-term trend of international oil prices, to build a large-scale model of usdcad, and build it with mathematical data from a macro perspective. There is a dual relationship between the surface and the inside called the trend structure in technical theory, that is, any trend in a large level must fall in a certain direction. The long-term trend is the general direction, and the time series modeling can be better in the changing market to capture the high and low points of our trading, so it will be more useful in higher frequency transactions.*
#
# *However, due to the limited frequency and range of data I can obtain, and in order to simplify the model, the data selected below are all daily frequencies.*

# ### b.	Causality & Statistical Signifiance
#
#    i. Utilize the Granger’s Causality Test to test foe causation in the variables that you have selected. Construct a matrix to show the p-values of the variables against one another (i.e., you also need to show that there are possible relationships between one another; not just USD/CAD).
#
#    ii. [Test for Statistical Significance Between Time Series] Use the Johassen’s test of cointegration to check statistical significance in your variables.
#
# ### c.	[Check for Stationarity] Implement an ADF test to check for stationarity. If it is not stationary, conduct a first-order differencing and re-check for stationarity. Show the ADF p-values for all the selected variables. If first-order differencing is not stationary, conduct a second-order differencing and re-run the ADF test.
#
# ### d.	[Model Selection] Using fit comparison estimates of AIC, BIC, FPE, and HQIC. Using the Fit comparison estimates to derive the most optimal number of lags for the model.
#

# ### Data Preprocessing
#

# In[1]:

# Data Preprocessing

import numpy as np
import pandas as pd
from datetime import datetime
import matplotlib.pyplot as plt

def Interpolation(data_process):

    """
    Interpolation

    Perform arithmetic interpolation on the missing values of a column of data
    """

    j = 0
    df_seq = []
    df = data_process.values

    # arithmetic interpolation, the value of 0 is also considered here.
    for i in range(len(df)):
        if df[i] != 0:
            diff = (df[i]-df[j])/(i-j)
            data_arange = np.arange(df[j], df[i], diff)
            j = i
            for num in range(len(data_arange)):
                df_seq.append(data_arange[num])

        #if i == len(df)-1 and df[i] != 0:
        #    df_seq.append(df[i])

        if i == len(df)-1 and df[i] == 0:
            end = df[j] + diff * (i-j)
            data_arange = np.arange(df[j], end, diff)
            for num in range(len(data_arange)):
                df_seq.append(data_arange[num])
            if len(df_seq) != len(df):
                df_seq.append(end)

    return df_seq

data = pd.read_csv('./HistoryData.csv')  # Read data.

data['Date'] = pd.to_datetime(data['Date'])
data.index = data['Date']  # Convert date format

# data.info()

data['WTI'] = data['WTI'].bfill() # There are few missing values in wti, so only mean interpolation is used.

data = data.fillna(0)

# These two are monthly data, and its trend needs to be retained, so use arithmetic interpolation
data['Sales_Rate'] = Interpolation(data['Sales_Rate'])
data['CrudeOilExports_UStoCanada'] = Interpolation(data['CrudeOilExports_UStoCanada'])
data = data.drop(data[(data['Sales_Rate']==0) |(data['CrudeOilExports_UStoCanada']==0)].index)

data.to_csv('./train_data.csv', index = False, header = True)
data.describe()

# In[4]:

# Data Visualization
DailyWTI = plt.figure(figsize = (10, 4))
ax = DailyWTI.add_subplot(111)
ax.set(title = 'WTI_DailyData',
        ylabel = 'Price', xlabel='Date')

plt.plot(data['Date'], data['WTI'])
plt.show()

DailyXAU_USD = plt.figure(figsize = (10, 4))
ax = DailyXAU_USD.add_subplot(111)
ax.set(title = 'XAU_USD_DailyData',
        ylabel = 'Price', xlabel='Date')

plt.plot(data['Date'], data['XAU_USD'])
plt.show()

DailySales_Rate = plt.figure(figsize = (10, 4))
ax = DailySales_Rate.add_subplot(111)
ax.set(title = 'Sales_Rate_DailyData',
        ylabel = 'Price', xlabel='Date')

plt.plot(data['Date'], data['Sales_Rate'])
plt.show()

DailyCrudeOilExports_UStoCanada = plt.figure(figsize = (10, 4))
ax = DailyCrudeOilExports_UStoCanada.add_subplot(111)
ax.set(title = 'CrudeOilExports_UStoCanada_DailyData',
        ylabel = 'Price', xlabel='Date')

plt.plot(data['Date'], data['CrudeOilExports_UStoCanada'])
plt.show()

# ### Granger Causality Tests
#
# ##### if p is less than 0.05, Granger causality is considered

# In[5]:

def Grangercausalitytests(data, maxlag_num):

    """
    Granger Causality Tests

    if p is less than 0.05, Granger causality is considered
    """

    from statsmodels.tsa.stattools import grangercausalitytests
    grangercausalitytests(data, maxlag=maxlag_num)

# All variables need to be tested
print('Grangercausalitytests Result between WTI and USD_CAD')
Grangercausalitytests(data[['WTI', 'USD_CAD']], 3)

print('\nGrangercausalitytests Result between XAU_USD and USD_CAD')
Grangercausalitytests(data[['XAU_USD', 'USD_CAD']], 3)

print('\nGrangercausalitytests Result between Sales_Rate and USD_CAD')
Grangercausalitytests(data[['Sales_Rate', 'USD_CAD']], 3)

print('\nGrangercausalitytests Result between CrudeOilExports_UStoCanada and USD_CAD')
Grangercausalitytests(data[['CrudeOilExports_UStoCanada', 'USD_CAD']], 3)

print('\nGrangercausalitytests Result between WTI and CrudeOilExports_UStoCanada')
Grangercausalitytests(data[['WTI', 'CrudeOilExports_UStoCanada']], 3)

print('\nGrangercausalitytests Result between Sales_Rate and CrudeOilExports_UStoCanada')
Grangercausalitytests(data[['Sales_Rate', 'CrudeOilExports_UStoCanada']], 3)

print('\nGrangercausalitytests Result between XAU_USD and CrudeOilExports_UStoCanada')
Grangercausalitytests(data[['XAU_USD', 'CrudeOilExports_UStoCanada']], 3)

print('\nGrangercausalitytests Result between WTI and Sales_Rate')
Grangercausalitytests(data[['WTI', 'Sales_Rate']], 3)

print('\nGrangercausalitytests Result between XAU_USD and Sales_Rate')
Grangercausalitytests(data[['XAU_USD', 'Sales_Rate']], 3)

print('\nGrangercausalitytests Result between WTI and XAU_USD')
Grangercausalitytests(data[['WTI', 'XAU_USD']], 3)

# ### ADF: Augmented Dickey-Fuller Unit Root  Tests
#
# ##### Check if it is a stationary series

# In[6]:

def ADF_diff(timeseries, name):

    """
    ADF: Augmented Dickey-Fuller unit root test.

    Regression: Constant and trend order to include {“c”,”ct”,”ctt”,”nc”}
    1. “c” : constant only (default).
    2. “ct” : constant and trend.
    3. “ctt” : constant, and linear and quadratic trend.
    4. “nc” : no constant, no trend.
    """

    import pandas as pd
    import statsmodels.api as sm
    import matplotlib.pyplot as plt
    from statsmodels.tsa.stattools import adfuller as ADF

    # Sequence after generating the differencing.
    timeseries_diff1 = timeseries.diff(1)
    timeseries_diff2 = timeseries_diff1.diff(1)

    timeseries_diff1 = timeseries_diff1.fillna(0)
    timeseries_diff2 = timeseries_diff2.fillna(0)

    # ADF unit root test -- ct
    print('Result of ADF--ct Test ')
    timeseries_adf = ADF(timeseries[name].tolist(), regression='ct')
    timeseries_diff1_adf = ADF(timeseries_diff1[name].tolist(), regression='ct')
    timeseries_diff2_adf = ADF(timeseries_diff2[name].tolist(), regression='ct')

    print('timeseries_adf : ', timeseries_adf)
    print('timeseries_diff1_adf : ', timeseries_diff1_adf)
    print('timeseries_diff2_adf : ', timeseries_diff2_adf)

    plt.figure(figsize=(12, 8))
    plt.plot(timeseries, label='Original', color='blue')
    plt.plot(timeseries_diff1, label='Diff1', color='red')
    plt.plot(timeseries_diff2, label='Diff2', color='purple')
    plt.legend(loc='best')
    plt.show()

    # ADF unit root test -- c
    print('Result of ADF--c Test ')
    timeseries_adf = ADF(timeseries[name].tolist(), regression='c')
    timeseries_diff1_adf = ADF(timeseries_diff1[name].tolist(), regression='c')
    timeseries_diff2_adf = ADF(timeseries_diff2[name].tolist(), regression='c')

    print('timeseries_adf : ', timeseries_adf)
    print('timeseries_diff1_adf : ', timeseries_diff1_adf)
    print('timeseries_diff2_adf : ', timeseries_diff2_adf)

    plt.figure(figsize=(12, 8))
    plt.plot(timeseries, label='Original', color='blue')
    plt.plot(timeseries_diff1, label='Diff1', color='red')
    plt.plot(timeseries_diff2, label='Diff2', color='purple')
    plt.legend(loc='best')
    plt.show()

    # ADF unit root test -- nc
    print('Result of ADF--nc Test ')
    timeseries_adf = ADF(timeseries[name].tolist(), regression='nc')
    timeseries_diff1_adf = ADF(timeseries_diff1[name].tolist(), regression='nc')
    timeseries_diff2_adf = ADF(timeseries_diff2[name].tolist(), regression='nc')

    print('timeseries_adf : ', timeseries_adf)
    print('timeseries_diff1_adf : ', timeseries_diff1_adf)
    print('timeseries_diff2_adf : ', timeseries_diff2_adf)

    plt.figure(figsize=(12, 8))
    plt.plot(timeseries, label='Original', color='blue')
    plt.plot(timeseries_diff1, label='Diff1', color='red')
    plt.plot(timeseries_diff2, label='Diff2', color='purple')
    plt.legend(loc='best')
    plt.show()

# Parse data with date for training
dateparse = lambda dates: datetime.strptime(dates, '%Y-%m-%d')
train_data = pd.read_csv('./train_data.csv', parse_dates=['Date'],
                                 index_col='Date', date_parser=dateparse)

# Extract the data of each variable for future use
data_WTI = pd.DataFrame(train_data['WTI'])
data_XAU_USD = pd.DataFrame(train_data['XAU_USD'])
data_Sales_Rate = pd.DataFrame(train_data['Sales_Rate'])
data_CrudeOilExports_UStoCanada = pd.DataFrame(train_data['CrudeOilExports_UStoCanada'])
data_USD_CAD = pd.DataFrame(train_data['USD_CAD'])

# ADF Tests
print("Resulty of ADF -- WTI")
ADF_diff(data_WTI, 'WTI')
print("Resulty of ADF -- XAU_USD")
ADF_diff(data_XAU_USD, 'XAU_USD')
print("Resulty of ADF -- Sales_Rate")
ADF_diff(data_Sales_Rate, 'Sales_Rate')
print("Resulty of ADF -- CrudeOilExports_UStoCanada")
ADF_diff(data_CrudeOilExports_UStoCanada, 'CrudeOilExports_UStoCanada')
print("Resulty of ADF -- USD_CAD")
ADF_diff(data_USD_CAD, 'USD_CAD')

# ### Fit Comparison Estimates of AIC, BIC, FPE, and HQIC.
#
# ##### Using fit comparison estimates of AIC, BIC, FPE, and HQIC.
# ##### Using the Fit comparison estimates to derive the most optimal number of lags for the model.

# In[7]:

def order_selection(timeseries):

    """
    Select the optimal lag order.

    Use Vector Auto Regression (VAR) Model to select order.
    The var model can pass a maximum number of lags and the order criterion to use for lag order selection.
    """

    import statsmodels.api as sm
    from statsmodels.tsa.api import VAR

    # Vector Auto Regression (VAR) Model
    var_model = VAR(timeseries)

    #Lag order selection
    order = var_model.select_order(10)
    print(order.summary())

    #var_results = var_model.fit(maxlags=5, ic='aic')
    #var_results.summary()

# Use original data to select order
order_selection(train_data)

# data after differencing
train_data_diff = train_data

data_WTI_diff1 = data_WTI.diff(1)
data_WTI_diff1 = data_WTI_diff1.fillna(0)
train_data_diff['WTI'] = data_WTI_diff1

data_XAU_USD_diff1 = data_XAU_USD.diff(1)
data_XAU_USD_diff1 = data_XAU_USD_diff1.fillna(0)
train_data_diff['XAU_USD'] = data_XAU_USD_diff1

data_Sales_Rate_diff1 = data_Sales_Rate.diff(1)
data_Sales_Rate_diff2 = data_Sales_Rate_diff1.diff(1)
data_Sales_Rate_diff2 = data_Sales_Rate_diff2.fillna(0)
train_data_diff['Sales_Rate'] = data_Sales_Rate_diff2

data_CrudeOilExports_UStoCanada_diff1 = data_CrudeOilExports_UStoCanada.diff(1)
data_CrudeOilExports_UStoCanada_diff2 = data_CrudeOilExports_UStoCanada_diff1.diff(1)
data_CrudeOilExports_UStoCanada_diff2 = data_CrudeOilExports_UStoCanada_diff2.fillna(0)
train_data_diff['CrudeOilExports_UStoCanada'] = data_CrudeOilExports_UStoCanada_diff2

data_USD_CAD_diff1 = data_USD_CAD.diff(1)
data_USD_CAD_diff1 = data_USD_CAD_diff1.fillna(0)
train_data_diff['USD_CAD'] = data_USD_CAD_diff1

# Use data after differencing to select order
order_selection(train_data_diff)

# ### Johassen’s test of cointegration

# In[8]:

def coint_johansen(timeseries):

    """
    Johassen’s test of cointegration

    Use the Johassen’s test of cointegration to check statistical significance in your variables.
    """

    from statsmodels.tsa.vector_ar import vecm

    # Johansen cointegration test, situation setting = continuous features
    jres1 = vecm.coint_johansen(timeseries, det_order=0, k_ar_diff=1)

    # View johansen cointegration test results
    johansen_result = vecm.select_coint_rank(timeseries,det_order=0, k_ar_diff=1, signif=0.1)
    print(johansen_result.summary())

    # Choice of rank
    print(johansen_result.rank)

    j_name=np.array([['Order'], ['trace statistics'], ['CV 90%'],  ['CV 95%'], ['CV 99%'] ] )
    print(j_name)

    j_order = jres1.ind
    print(j_order) # Order of eigenvalues

    j_lr1 = jres1.lr1
    print(j_lr1) # trace statistics

    j_cvt = jres1.cvt
    print(j_cvt) # Critical values (90%, 95%, 99%) of trace statistic

    # j_result = np.vstack(( j_order,  j_lr1, j_cvt))
    # print(j_result)

    pd.DataFrame(jres1.r0t).plot(kind='bar')  # Residuals for Δ

coint_johansen(train_data_diff)

# ### e.	[Forecast] Utilizing your model, forecast the daily close for the next five days, include a 75% confidence interval. Derive the USD/CAD closing price range for the next 5 days (i.e., Min & Max).
#

# ### VECM Model

# In[346]:

# Use VECM Model to Forecast
mod = vecm.VECM(train_data_diff, k_ar_diff=1, coint_rank=5, freq='B', deterministic="ci")
res = mod.fit()
print(res.summary())

# In[347]:

"""
Notes:

forecast - ndarray (steps x neqs) or three ndarrays

In case of a point forecast: each row of the returned ndarray represents the forecast of the neqs variables for a specific period.
The first row (index [0]) is the forecast for the next period, the last row (index [steps-1]) is the steps-periods-ahead- forecast.
"""

# An image of data that predicts one step forward (showing confidence interval, confidence interval 75%)
res.plot_forecast(steps=5,plot_conf_int=True,alpha=0.75)

# Predict the value one step forward, with 75% confidence
res.predict(steps=5, alpha=0.75)

# In[ ]:

# In[ ]:

import numpy,import statsmodels,from statsmodels


# Copyright 2019 Xanadu Quantum Technologies Inc.

# 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.
r"""
Python library
==============

.. currentmodule:: thewalrus

This is the top level module of the The Walrus Python interface,
containing functions for computing the hafnian, loop hafnian,
and torontonian of matrices.

Algorithm terminology
---------------------

Eigenvalue hafnian algorithm
    The algorithm described in
    *A faster hafnian formula for complex matrices and its benchmarking on a supercomputer*,
    :cite:`bjorklund2018faster`.
    This algorithm scales like :math:`\mathcal{O}(n^3 2^{n/2})`, and supports calculation of
    the loop hafnian.

Recursive hafnian algorithm
    The algorithm described in *Counting perfect matchings as fast as Ryser* :cite:`bjorklund2012counting`.
    This algorithm scales like :math:`\mathcal{O}(n^4 2^{n/2})`. This algorithm does not
    currently support the loop hafnian.

Repeating hafnian algorithm
    The algorithm described in *From moments of sum to moments of product*, :cite:`kan2008moments`.
    This method is more efficient for matrices with repeated rows and columns, and supports caclulation of
    the loop hafnian.

Approximate hafnian algorithm
    The algorithm described in *Polynomial time algorithms to approximate permanents and mixed discriminants
    within a simply exponential factor*, :cite:`barvinok1999polynomial`.
    This algorithm allows us to efficiently approximate the hafnian of
    matrices with non-negative elements. This is done by sampling determinants;
    the larger the number of samples taken, the higher the accuracy.

Batched hafnian algorithm
    An algorithm that allows the calculation of hafnians of all reductions of
    a given matrix up to the cutoff (resolution) provided. Internally, this algorithm
    makes use of the multidimensional Hermite polynomials as per
    *The calculation of multidimensional Hermite polynomials and Gram-Charlier coefficients*
    :cite:`berkowitz1970calculation`.

Low-rank hafnian algorithm
    An algorithm that allows to calculate the hafnian of an :math:`r`-rank matrix :math:`\bm{A}` of size :math:`n \times n`
    by factorizing it as :math:`\bm{A} = \bm{G} \bm{G}^T` where :math:`\bm{G}` is of size :math:`n \times r`. The algorithm
    is described in Appendix C of
    *A faster hafnian formula for complex matrices and its benchmarking on a supercomputer*,
    :cite:`bjorklund2018faster`.

Python wrappers
---------------

.. autosummary::
    hafnian
    hafnian_repeated
    hafnian_batched
    tor
    perm
    permanent_repeated
    hermite_multidimensional

Pure Python functions
---------------------

.. autosummary::
    reduction
    version
    low_rank_hafnian
"""
# pylint: disable=wrong-import-position
import os
import platform

import numpy as np

import thewalrus.quantum

from ._hafnian import (
    haf_complex,
    haf_int,
    haf_real,
    haf_rpt_complex,
    haf_rpt_real,
    hafnian,
    hafnian_repeated,
    reduction,
)
from ._low_rank_haf import low_rank_hafnian
from ._hermite_multidimensional import hafnian_batched, hermite_multidimensional
from ._permanent import perm, perm_complex, perm_real, permanent_repeated
from ._torontonian import tor, threshold_detection_prob_displacement, threshold_detection_prob, numba_tor
from ._version import __version__

__all__ = [
    "hafnian",
    "hafnian_repeated",
    "hafnian_batched",
    "tor",
    "perm",
    "permanent_repeated",
    "reduction",
    "hermite_multidimensional",
    "version",
]

def version():
    r"""
    Get version number of The Walrus

    Returns:
      str: The package version number
    """
    return __version__


__author__ = 'Robert Meyer'

import numpy as np
import inspect
import os # For path names being viable under Windows and Linux

from pypet import Environment, Parameter, ArrayParameter, Trajectory
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Here we will see how we can write our own custom parameters and how we can use
# it with a trajectory.

# Now we want to do a more sophisticated simulations, we will integrate a differential equation
# with an Euler scheme

# Let's first define our job to do
def euler_scheme(traj, diff_func):
    """Simulation function for Euler integration.

    :param traj:

        Container for parameters and results

    :param diff_func:

        The differential equation we want to integrate

    """

    steps = traj.steps
    initial_conditions = traj.initial_conditions
    dimension = len(initial_conditions)

    # This array will collect the results
    result_array = np.zeros((steps,dimension))
    # Get the function parameters stored into `traj` as a dictionary
    # with the (short) names as keys :
    func_params_dict = traj.func_params.f_to_dict(short_names=True, fast_access=True)
    # Take initial conditions as first result
    result_array[0] = initial_conditions

    # Now we compute the Euler Scheme steps-1 times
    for idx in range(1,steps):
        result_array[idx] = diff_func(result_array[idx-1], **func_params_dict) * traj.dt + \
                            result_array[idx-1]
    # Note the **func_params_dict unzips the dictionary, it's the reverse of **kwargs in function
    # definitions!

    #Finally we want to keep the results
    traj.f_add_result('euler_evolution', data=result_array, comment='Our time series data!')

# Ok, now we want to make our own (derived) parameter that stores source code of python functions.
# We do NOT want a parameter that stores an executable function. This would complicate
# the problem a lot. If you have something like that in mind, you might wanna take a look
# at the marshal (http://docs.python.org/2/library/marshal) module
# or dill (https://pypi.python.org/pypi/dill) package.
# Our intention here is to define a parameter that we later on use as a derived parameter
# to simply keep track of the source code we use ('git' would be, of course, the better solution
# but this is just an illustrative example)
class FunctionParameter(Parameter):
    # We need to override the `f_set` function and simply extract the the source code if our
    # item is callable and store this instead.
    def f_set(self, data):
        if callable(data):
            data = inspect.getsource(data)
        return super(FunctionParameter, self).f_set(data)

    # For more complicate parameters you might consider implementing:
    # `f_supports` (we do not need it since we convert the data to stuff the parameter already
    #    supports, and that is strings!)
    #
    # and
    # the private functions
    #
    # `_values_of_same_type` (to tell whether data is similar, i.e. of two data items agree in their
    #   type, this is important to only allow exploration within the same dimension.
    #   For instance, a parameter that stores integers, should only explore integers etc.)
    #
    # and
    #
    # `_equal_values` (to tell if two data items are equal. This is important for merging if you
    #       want to erase duplicate parameter points. The trajectory needs to know when a
    #       parameter space point was visited before.)
    #
    # and
    #
    # `_store` (to be able to turn complex data into basic types understood by the storage service)
    #
    # and
    #
    # `_load` (to be able to recover your complex data form the basic types understood by the storage
    # service)
    #
    # But for now we will rely on the parent functions and hope for the best!

# Ok now let's follow the ideas in the final section of the cookbook and let's
# have a part in our simulation that only defines the parameters.
def add_parameters(traj):
    """Adds all necessary parameters to the `traj` container"""

    traj.f_add_parameter('steps', 10000, comment='Number of time steps to simulate')
    traj.f_add_parameter('dt', 0.01, comment='Step size')

    # Here we want to add the initial conditions as an array parameter. We will simulate
    # a 3-D differential equation, the Lorenz attractor.
    traj.f_add_parameter(ArrayParameter,'initial_conditions', np.array([0.0,0.0,0.0]),
                         comment = 'Our initial conditions, as default we will start from'
                                   ' origin!')

    # We will group all parameters of the Lorenz differential equation into the group 'func_params'
    traj.f_add_parameter('func_params.sigma', 10.0)
    traj.f_add_parameter('func_params.beta', 8.0/3.0)
    traj.f_add_parameter('func_params.rho', 28.0)

    #For the fun of it we will annotate the  group
    traj.func_params.v_annotations.info='This group contains as default the original values chosen ' \
                                   'by Edward Lorenz in 1963. Check it out on wikipedia ' \
                                   '(https://en.wikipedia.org/wiki/Lorenz_attractor)!'

# We need to define the lorenz function, we will assume that the value array is 3 dimensional,
# First dimension contains the x-component, second y-component, and third the z-component
def diff_lorenz(value_array, sigma, beta, rho):
    """The Lorenz attractor differential equation

    :param value_array: 3d array containing the x,y, and z component values.
    :param sigma: Constant attractor parameter
    :param beta: FConstant attractor parameter
    :param rho: Constant attractor parameter

    :return: 3d array of the Lorenz system evaluated at `value_array`

    """
    diff_array = np.zeros(3)
    diff_array[0] = sigma * (value_array[1]-value_array[0])
    diff_array[1] = value_array[0] * (rho - value_array[2]) - value_array[1]
    diff_array[2] = value_array[0] * value_array[1] - beta * value_array[2]

    return diff_array

# And here goes our main function
def main():

    filename = os.path.join('hdf5', 'example_05.hdf5')
    env = Environment(trajectory='Example_05_Euler_Integration',
                      filename=filename,
                      file_title='Example_05_Euler_Integration',
                      overwrite_file=True,
                      comment='Go for Euler!')

    traj = env.trajectory
    trajectory_name = traj.v_name

    # 1st a) phase parameter addition
    add_parameters(traj)

    # 1st b) phase preparation
    # We will add the differential equation (well, its source code only) as a derived parameter
    traj.f_add_derived_parameter(FunctionParameter,'diff_eq', diff_lorenz,
                                 comment='Source code of our equation!')

    # We want to explore some initial conditions
    traj.f_explore({'initial_conditions' : [
        np.array([0.01,0.01,0.01]),
        np.array([2.02,0.02,0.02]),
        np.array([42.0,4.2,0.42])
    ]})
    # 3 different conditions are enough for an illustrative example

    # 2nd phase let's run the experiment
    # We pass `euler_scheme` as our top-level simulation function and
    # the Lorenz equation 'diff_lorenz' as an additional argument
    env.run(euler_scheme, diff_lorenz)

    # We don't have a 3rd phase of post-processing here

    # 4th phase analysis.
    # I would recommend to do post-processing completely independent from the simulation,
    # but for simplicity let's do it here.

    # Let's assume that we start all over again and load the entire trajectory new.
    # Yet, there is an error within this approach, do you spot it?
    del traj
    traj = Trajectory(filename=filename)

    # We will only fully load parameters and derived parameters.
    # Results will be loaded manually later on.
    try:
        # However, this will fail because our trajectory does not know how to
        # build the FunctionParameter. You have seen this coming, right?
        traj.f_load(name=trajectory_name, load_parameters=2, load_derived_parameters=2,
                    load_results=1)
    except ImportError as e:

        print('That did\'nt work, I am sorry: %s ' % str(e))

        # Ok, let's try again but this time with adding our parameter to the imports
        traj = Trajectory(filename=filename,
                           dynamically_imported_classes=FunctionParameter)

        # Now it works:
        traj.f_load(name=trajectory_name, load_parameters=2, load_derived_parameters=2,
                    load_results=1)

    #For the fun of it, let's print the source code
    print('\n ---------- The source code of your function ---------- \n %s' % traj.diff_eq)

    # Let's get the exploration array:
    initial_conditions_exploration_array = traj.f_get('initial_conditions').f_get_range()
    # Now let's plot our simulated equations for the different initial conditions:
    # We will iterate through the run names
    for idx, run_name in enumerate(traj.f_get_run_names()):

        #Get the result of run idx from the trajectory
        euler_result = traj.results.f_get(run_name).euler_evolution
        # Now we manually need to load the result. Actually the results are not so large and we
        # could load them all at once. But for demonstration we do as if they were huge:
        traj.f_load_item(euler_result)
        euler_data = euler_result.data

        #Plot fancy 3d plot
        fig = plt.figure(idx)
        ax = fig.gca(projection='3d')
        x = euler_data[:,0]
        y = euler_data[:,1]
        z = euler_data[:,2]
        ax.plot(x, y, z, label='Initial Conditions: %s' % str(initial_conditions_exploration_array[idx]))
        plt.legend()
        plt.show()

        # Now we free the data again (because we assume its huuuuuuge):
        del euler_data
        euler_result.f_empty()

    # You have to click through the images to stop the example_05 module!

    # Finally disable logging and close all log-files
    env.disable_logging()

if __name__ == '__main__':
    main()


#!/usr/bin/env python
u"""
grace_date.py
Written by Tyler Sutterley (05/2021)

Reads index file from podaac_grace_sync.py or gfz_isdc_grace_ftp.py
Parses dates of each GRACE/GRACE-FO file and assigns the month number
Creates an index of dates for GRACE/GRACE-FO files

INPUTS:
    base_dir: Working data directory for GRACE/GRACE-FO data

OPTIONS:
    PROC: GRACE data processing center (CSR/CNES/JPL/GFZ/GRAZ/COSTG/SWARM)
    DREL: GRACE/GRACE-FO Data Release (RL03 for CNES) (RL06 for CSR/GFZ/JPL)
    DSET: GRACE dataset (GAA/GAB/GAC/GAD/GSM)
        GAA is the non-tidal atmospheric correction
        GAB is the non-tidal oceanic correction
        GAC is the combined non-tidal atmospheric and oceanic correction
        GAD is the GRACE ocean bottom pressure product
        GSM is corrected monthly GRACE/GRACE-FO static field product
    OUTPUT: create index of dates for GRACE/GRACE-FO data
    MODE: permissions mode of output files

OUTPUTS:
    dictionary of GRACE/GRACE-FO files indexed by month

PYTHON DEPENDENCIES:
    numpy: Scientific Computing Tools For Python
        https://numpy.org
        https://numpy.org/doc/stable/user/numpy-for-matlab-users.html
    dateutil: powerful extensions to datetime
        https://dateutil.readthedocs.io/en/stable/
    future: Compatibility layer between Python 2 and Python 3
        https://python-future.org/

PROGRAM DEPENDENCIES:
    time.py: utilities for calculating time operations

UPDATE HISTORY:
    Updated 05/2021: define int/float precision to prevent deprecation warning
    Updated 02/2021: use adjust_months function to fix special months cases
    Updated 12/2020: Add SWARM data compilance
    Updated 12/2020: using utilities from time module
    Updated 11/2020: updated for CNES RL04 & RL05 and GRAZ 2018 (monthly fields)
    Updated 10/2020: use argparse to set command line parameters
    Updated 07/2020: added function docstrings
    Updated 03/2020: for public release
    Updated 11/2018: updated regular expression pattern for RL06 GFZ
    Updated 08/2018: using full release string (RL05 instead of 5)
    Updated 06/2018: using python3 compatible octal and input
    Updated 05/2018: can read new GRACE file format for RL06 and GRACE-FO
    Updated 10/2017: added option OUTPUT to write the text file with GRACE dates
        now will output from the function the GRACE month and file name
        adjusted regular expression for extracting parameters from filename
    Updated 02/2017: added mode to modify the permissions of the output file
    Updated 05-06/2016: using __future__ print function, format date lines
    Updated 01/2016: minor clean up.  using enumerate for loop
    Updated 10/2015: added manual fix for month 161 (centered in 160)
    Updated 05/2015: additional 2 digits to date file to reduce the differences
        in dates if importing from date file or binary data file
    Updated 03/2015: added main definition for running from command line
        further generalization, calculates the Julian date of the GRACE date
    Updated 11/2014: output more decimal points for date
    Updated 09/2014: using regular expressions, general code updates
    Updated 03/2014: printing GRACE month with zero-padding
    Updated 02/2014: minor update to if statements
    Updated 01/2014: updated for CNES RL03 (monthly fields)
    Updated 10/2013: created a directory with both RL04 and RL05 (combining)
        made a slight edit for the drift rates
    Updated 09/2013: changed dating schemes for all products
        for CNES: Solution 1 == 001, versus 10-day from start of 2002
    Updated 05/2013: modified for use with GUI program
    Updated 07/2012: changed some variable names, and saving variables
        start_yr, start_day, end_yr, end_day
    Updated 07/2012: fix missing date for CSR RL05
        One of the months 119 registered as 118 as half of 119 is in 118 due to
        accelerometer issues during 119
    Updated 06/2012: fixes missing dates to be automatic
        Also added options to enter the processing center, the data release and
        the dataset from an external main level program
    Updated 04/2012: changes for RL05 data
"""
from __future__ import print_function

import os
import re
import argparse
import numpy as np
import gravity_toolkit.time

def grace_date(base_dir, PROC='', DREL='', DSET='', OUTPUT=True, MODE=0o775):
    """
    Reads index file from podaac_grace_sync.py or gfz_isdc_grace_ftp.py
    Parses dates of each GRACE/GRACE-FO file and assigns the month number
    Creates an index of dates for GRACE/GRACE-FO files

    Arguments
    ---------
    base_dir: working data directory

    Keyword arguments
    -----------------
    PROC: GRACE data processing center
        CSR: University of Texas Center for Space Research
        GFZ: German Research Centre for Geosciences (GeoForschungsZentrum)
        JPL: Jet Propulsion Laboratory
        CNES: French Centre National D'Etudes Spatiales
        GRAZ: Institute of Geodesy from GRAZ University of Technology
        COSTG: International Combination Service for Time-variable Gravity Fields

        SWARM: gravity data from SWARM satellite
    DREL: GRACE/GRACE-FO data release
    DSET: GRACE/GRACE-FO dataset
        GAA: non-tidal atmospheric correction
        GAB: non-tidal oceanic correction
        GAC: combined non-tidal atmospheric and oceanic correction
        GAD: ocean bottom pressure product
        GSM: corrected monthly static gravity field product
    OUTPUT: create index file of dates for GRACE/GRACE-FO data
    MODE: Permission mode of directories and files

    Returns
    -------
    output_files: dictionary of GRACE/GRACE-FO files indexed by month
    """

    #--  Directory of exact product
    grace_dir = os.path.join(base_dir, PROC, DREL, DSET)
    #-- input index file containing GRACE data filenames
    with open(os.path.join(grace_dir, 'index.txt'),'r') as f:
        input_files = f.read().splitlines()

    #--  number of lines in input_files
    n_files = len(input_files)

    #-- define date variables
    start_yr = np.zeros((n_files))#-- year start date
    end_yr = np.zeros((n_files))#-- year end date
    start_day = np.zeros((n_files))#-- day number start date
    end_day = np.zeros((n_files))#-- day number end date
    mid_day = np.zeros((n_files))#-- mid-month day
    tot_days = np.zeros((n_files))#-- number of days since Jan 2002
    tdec = np.zeros((n_files))#-- tdec is the date in decimal form
    mon = np.zeros((n_files,),dtype=np.int64)#-- GRACE/GRACE-FO month number

    if PROC in ('CSR', 'GFZ', 'JPL', 'CNES', 'COSTG'):
        #-- compile numerical expression operator for parameters from files
        #-- will work with previous releases and releases for GRACE-FO
        #-- UTCSR: The University of Texas at Austin Center for Space Research
        #-- EIGEN: GFZ German Research Center for Geosciences (RL01-RL05)
        #-- GFZOP: GFZ German Research Center for Geosciences (RL06+GRACE-FO)
        #-- JPLEM: NASA Jet Propulsion Laboratory (harmonic solutions)
        #-- JPLMSC: NASA Jet Propulsion Laboratory (mascon solutions)
        #-- GRGS: CNES Groupe de Recherche de Géodésie Spatiale
        regex_pattern = (r'(.*?)-2_(\d+)-(\d+)_(.*?)_({0})_(.*?)_(\d+)(.*?)'
           r'(\.gz|\.gfc|\.txt)?$').format(r'UTCSR|EIGEN|GFZOP|JPLEM|JPLMSC|GRGS|COSTG')
    elif PROC == 'GRAZ':
        # -- GRAZ: Institute of Geodesy from GRAZ University of Technology
        regex_pattern = (r'(.*?)-({0})_(.*?)_(\d+)-(\d+)'
                         r'(\.gz|\.gfc|\.txt)').format(r'Grace_operational|Grace2018')
    elif PROC == 'SWARM':
        # -- SWARM: data from SWARM satellite
        regex_pattern = (r'({0})_(.*?)_(EGF_SHA_2)__(.*?)_(.*?)_(.*?)'
                         r'(\.gz|\.gfc|\.txt)').format(r'SW')
    else:
        raise ValueError("Unknown PROC value:", PROC)

    rx = re.compile(regex_pattern, re.VERBOSE)

    #-- for each data file
    for t, infile in enumerate(input_files):
        #-- extract parameters from input filename
        if PROC in ('CSR', 'GFZ', 'JPL', 'CNES', 'COSTG'):
            PFX,start_date,end_date,AUX,PRC,F1,DRL,F2,SFX = rx.findall(infile).pop()

            #-- find start date, end date and number of days
            start_yr[t] = np.float64(start_date[:4])
            end_yr[t] = np.float64(end_date[:4])
            start_day[t] = np.float64(start_date[4:])
            end_day[t] = np.float64(end_date[4:])

        elif PROC == 'GRAZ' or PROC == 'SWARM':
            if PROC == 'GRAZ':
                PFX,SAT,trunc,year,month,SFX = rx.findall(infile).pop()
                #-- find start year, end year
                start_yr[t] = np.float(year)
                end_yr[t] = np.float(year)
            elif PROC == 'SWARM':
                SAT, tmp, PROD, start_date, end_date, RL, SFX = rx.findall(os.path.basename(infile)).pop()

                start_yr[t] = int(start_date[:4])
                end_yr[t] = int(end_date[:4])
                month = int(start_date[4:6])

            #-- Calculation of total days since start of campaign
            #-- Get information on the current year (day per month and day per year)
            dpm = gravity_toolkit.time.dpm_count(start_yr[t])

            #-- find start day, end day
            start_day[t] = np.sum(dpm[:np.int(month) - 1]) + 1
            end_day[t] = np.sum(dpm[:np.int(month)])

        # -- number of days in the starting year for leap and standard years
        dpy = gravity_toolkit.time.calendar_days(start_yr[t]).sum()
        #-- end date taking into account measurements taken on different years
        end_cyclic = (end_yr[t]-start_yr[t])*dpy + end_day[t]
        #-- calculate mid-month value
        mid_day[t] = np.mean([start_day[t], end_cyclic])

        #-- calculate Modified Julian Day from start_yr and mid_day
        MJD = gravity_toolkit.time.convert_calendar_dates(start_yr[t],
            1.0,mid_day[t],epoch=(1858,11,17,0,0,0))
        #-- convert from Modified Julian Days to calendar dates
        cal_date = gravity_toolkit.time.convert_julian(MJD+2400000.5)

        #-- Calculating the mid-month date in decimal form
        tdec[t] = start_yr[t] + mid_day[t]/dpy

        #-- Calculation of total days since start of campaign
        count = 0
        n_yrs = np.int64(start_yr[t]-2002)
        #-- for each of the GRACE years up to the file year
        for iyr in range(n_yrs):
            #-- year
            year = 2002 + iyr
            #-- add all days from prior years to count
            #-- number of days in year i (if leap year or standard year)
            count += gravity_toolkit.time.calendar_days(year).sum()

        #-- calculating the total number of days since 2002
        tot_days[t] = np.mean([count+start_day[t], count+end_cyclic])

        #-- Calculates the month number (or 10-day number for CNES RL01,RL02)
        if ((PROC == 'CNES') and (DREL in ('RL01','RL02'))):
            mon[t] = np.round(1.0+(tot_days[t]-tot_days[0])/10.0)
        else:
            #-- calculate the GRACE/GRACE-FO month (Apr02 == 004)
            #-- https://grace.jpl.nasa.gov/data/grace-months/
            #-- Notes on special months (e.g. 119, 120) below
            mon[t] = 12*(cal_date['year']-2002) + cal_date['month']

    #-- The 'Special Months' (Nov 2011, Dec 2011 and April 2012) with
    #-- Accelerometer shutoffs make the relation between month number
    #-- and date more complicated as days from other months are used
    #-- For CSR and GFZ: Nov 2011 (119) is centered in Oct 2011 (118)
    #-- For JPL: Dec 2011 (120) is centered in Jan 2012 (121)
    #-- For all: May 2015 (161) is centered in Apr 2015 (160)
    mon = gravity_toolkit.time.adjust_months(mon)

    #-- Output GRACE/GRACE-FO date ascii file
    if OUTPUT:
        date_file = '{0}_{1}_DATES.txt'.format(PROC, DREL)
        fid = open(os.path.join(grace_dir,date_file), 'w')
        #-- date file header information
        args = ('Mid-date','Month','Start_Day','End_Day','Total_Days')
        print('{0} {1:>10} {2:>11} {3:>10} {4:>13}'.format(*args),file=fid)

    #-- create python dictionary mapping input file names with GRACE months
    grace_files = {}
    #-- for each data file
    for t, infile in enumerate(input_files):
        #-- add file to python dictionary mapped to GRACE/GRACE-FO month
        grace_files[mon[t]] = os.path.join(grace_dir,infile)
        #-- print to GRACE dates ascii file (NOTE: tot_days will be rounded)
        if OUTPUT:
            print(('{0:13.8f} {1:03d} {2:8.0f} {3:03.0f} {4:8.0f} {5:03.0f} '
                '{6:8.0f}').format(tdec[t],mon[t],start_yr[t],start_day[t],
                end_yr[t],end_day[t],tot_days[t]), file=fid)

    #-- close date file
    #-- set permissions level of output date file
    if OUTPUT:
        fid.close()
        os.chmod(os.path.join(grace_dir, date_file), MODE)

    #-- return the python dictionary that maps GRACE months with GRACE files
    return grace_files

#-- PURPOSE: program that calls grace_date() with set parameters
def main():
    #-- command line parameters
    #-- Read the system arguments listed after the program
    parser = argparse.ArgumentParser(
        description="""Parses dates of each GRACE/GRACE-FO file and
            assigns the month number.
            Creates an index of dates for GRACE/GRACE-FO files.
            """
    )
    #-- working data directory
    parser.add_argument('--directory','-D',
        type=lambda p: os.path.abspath(os.path.expanduser(p)),
        default=os.getcwd(),
        help='Working data directory')
    #-- GRACE/GRACE-FO data processing center
    parser.add_argument('--center','-c',
        metavar='PROC', type=str, nargs='+',
        default=['CSR','GFZ','JPL'],
        choices=['CSR','GFZ','JPL', 'CNES','GRAZ','SWARM', 'COSTG'],
        help='GRACE/GRACE-FO Processing Center')
    #-- GRACE/GRACE-FO data release
    parser.add_argument('--release','-r',
        metavar='DREL', type=str, nargs='+',
        default=['RL06'],
        help='GRACE/GRACE-FO Data Release')
    #-- GRACE/GRACE-FO data product
    parser.add_argument('--product','-p',
        metavar='DSET', type=str.upper, nargs='+',
        default=['GAC','GAD','GSM'],
        choices=['GAA','GAB','GAC','GAD','GSM'],
        help='GRACE/GRACE-FO Level-2 data product')
    #-- output GRACE/GRACE-FO ascii date file
    parser.add_argument('--output','-O',
        default=False, action='store_true',
        help='Overwrite existing data')
    #-- permissions mode of the local directories and files (number in octal)
    parser.add_argument('--mode','-M',
        type=lambda x: int(x,base=8), default=0o775,
        help='permissions mode of output files')
    args,_ = parser.parse_known_args()

    #-- run GRACE/GRACE-FO date program
    for pr in args.center:
        for rl in args.release:
            for ds in args.product:
                grace_date(args.directory, PROC=pr, DREL=rl, DSET=ds,
                    OUTPUT=args.output, MODE=args.mode)

#-- run main program
if __name__ == '__main__':
    main()