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

# # Project 3: Smart Beta Portfolio and Portfolio Optimization
#
# ## Overview
#
#
# Smart beta has a broad meaning, but we can say in practice that when we use the universe of stocks from an index, and then apply some weighting scheme other than market cap weighting, it can be considered a type of smart beta fund.  A Smart Beta portfolio generally gives investors exposure or "beta" to one or more types of market characteristics (or factors) that are believed to predict prices while giving investors a diversified broad exposure to a particular market. Smart Beta portfolios generally target momentum, earnings quality, low volatility, and dividends or some combination. Smart Beta Portfolios are generally rebalanced infrequently and follow relatively simple rules or algorithms that are passively managed.  Model changes to these types of funds are also rare requiring prospectus filings with US Security and Exchange Commission in the case of US focused mutual funds or ETFs.. Smart Beta portfolios are generally long-only, they do not short stocks.
#
# In contrast, a purely alpha-focused quantitative fund may use multiple models or algorithms to create a portfolio. The portfolio manager retains discretion in upgrading or changing the types of models and how often to rebalance the portfolio in attempt to maximize performance in comparison to a stock benchmark.  Managers may have discretion to short stocks in portfolios.
#
# Imagine you're a portfolio manager, and wish to try out some different portfolio weighting methods.
#
# One way to design portfolio is to look at certain accounting measures (fundamentals) that, based on past trends, indicate stocks that produce better results.
#
#
# For instance, you may start with a hypothesis that dividend-issuing stocks tend to perform better than stocks that do not. This may not always be true of all companies; for instance, Apple does not issue dividends, but has had good historical performance.  The hypothesis about dividend-paying stocks may go something like this:
#
# Companies that regularly issue dividends may also be more prudent in allocating their available cash, and may indicate that they are more conscious of prioritizing shareholder interests.  For example, a CEO may decide to reinvest cash into pet projects that produce low returns.  Or, the CEO may do some analysis, identify that reinvesting within the company produces lower returns compared to a diversified portfolio, and so decide that shareholders would be better served if they were given the cash (in the form of dividends).  So according to this hypothesis, dividends may be both a proxy for how the company is doing (in terms of earnings and cash flow), but also a signal that the company acts in the best interest of its shareholders.  Of course, it's important to test whether this works in practice.
#
#
# You may also have another hypothesis, with which you wish to design a portfolio that can then be made into an ETF.  You may find that investors may wish to invest in passive beta funds, but wish to have less risk exposure (less volatility) in their investments.  The goal of having a low volatility fund that still produces returns similar to an index may be appealing to investors who have a shorter investment time horizon, and so are more risk averse.
#
# So the objective of your proposed portfolio is to design a portfolio that closely tracks an index, while also minimizing the portfolio variance.  Also, if this portfolio can match the returns of the index with less volatility, then it has a higher risk-adjusted return (same return, lower volatility).
#
# Smart Beta ETFs can be designed with both of these two general methods (among others): alternative weighting and minimum volatility ETF.
#
#
# ## Instructions
# Each problem consists of a function to implement and instructions on how to implement the function.  The parts of the function that need to be implemented are marked with a `# TODO` comment. After implementing the function, run the cell to test it against the unit tests we've provided. For each problem, we provide one or more unit tests from our `project_tests` package. These unit tests won't tell you if your answer is correct, but will warn you of any major errors. Your code will be checked for the correct solution when you submit it to Udacity.
#
# ## Packages
# When you implement the functions, you'll only need to you use the packages you've used in the classroom, like [Pandas](https://pandas.pydata.org/) and [Numpy](http://www.numpy.org/). These packages will be imported for you. We recommend you don't add any import statements, otherwise the grader might not be able to run your code.
#
# The other packages that we're importing are `helper`, `project_helper`, and `project_tests`. These are custom packages built to help you solve the problems.  The `helper` and `project_helper` module contains utility functions and graph functions. The `project_tests` contains the unit tests for all the problems.
# ### Install Packages

# In[1]:

import sys
get_ipython().system('{sys.executable} -m pip install -r requirements.txt')

# In[2]:

import pandas as pd
import numpy as np
import helper
import project_helper
import project_tests

#
# ### Load Packages

# ## Market Data
# ### Load Data
# For this universe of stocks, we'll be selecting large dollar volume stocks. We're using this universe, since it is highly liquid.

# In[3]:

df = pd.read_csv('../../data/project_3/eod-quotemedia.csv')

percent_top_dollar = 0.2
high_volume_symbols = project_helper.large_dollar_volume_stocks(df, 'adj_close', 'adj_volume', percent_top_dollar)
df = df[df['ticker'].isin(high_volume_symbols)]

close = df.reset_index().pivot(index='date', columns='ticker', values='adj_close')
volume = df.reset_index().pivot(index='date', columns='ticker', values='adj_volume')
dividends = df.reset_index().pivot(index='date', columns='ticker', values='dividends')

# ### View Data
# To see what one of these 2-d matrices looks like, let's take a look at the closing prices matrix.

# In[4]:

project_helper.print_dataframe(close)

# # Part 1: Smart Beta Portfolio
# In Part 1 of this project, you'll build a portfolio using dividend yield to choose the portfolio weights. A portfolio such as this could be incorporated into a smart beta ETF.  You'll compare this portfolio to a market cap weighted index to see how well it performs.
#
# Note that in practice, you'll probably get the index weights from a data vendor (such as companies that create indices, like MSCI, FTSE, Standard and Poor's), but for this exercise we will simulate a market cap weighted index.
#
# ## Index Weights
# The index we'll be using is based on large dollar volume stocks. Implement `generate_dollar_volume_weights` to generate the weights for this index. For each date, generate the weights based on dollar volume traded for that date. For example, assume the following is close prices and volume data:
# ```
#                  Prices
#                A         B         ...
# 2013-07-08     2         2         ...
# 2013-07-09     5         6         ...
# 2013-07-10     1         2         ...
# 2013-07-11     6         5         ...
# ...            ...       ...       ...
#
#                  Volume
#                A         B         ...
# 2013-07-08     100       340       ...
# 2013-07-09     240       220       ...
# 2013-07-10     120       500       ...
# 2013-07-11     10        100       ...
# ...            ...       ...       ...
# ```
# The weights created from the function `generate_dollar_volume_weights` should be the following:
# ```
#                A         B         ...
# 2013-07-08     0.126..   0.194..   ...
# 2013-07-09     0.759..   0.377..   ...
# 2013-07-10     0.075..   0.285..   ...
# 2013-07-11     0.037..   0.142..   ...
# ...            ...       ...       ...
# ```

# In[5]:

def generate_dollar_volume_weights(close, volume):
    """
    Generate dollar volume weights.

    Parameters
    ----------
    close : DataFrame
        Close price for each ticker and date
    volume : str
        Volume for each ticker and date

    Returns
    -------
    dollar_volume_weights : DataFrame
        The dollar volume weights for each ticker and date
    """

    product = close*volume

    weights=product.apply(lambda r : r/sum(r),axis=1)

    assert close.index.equals(volume.index)
    assert close.columns.equals(volume.columns)

    #TODO: Implement function

    return weights

project_tests.test_generate_dollar_volume_weights(generate_dollar_volume_weights)

# ### View Data
# Let's generate the index weights using `generate_dollar_volume_weights` and view them using a heatmap.

# In[6]:

index_weights = generate_dollar_volume_weights(close, volume)
project_helper.plot_weights(index_weights, 'Index Weights')

# ## Portfolio Weights
# Now that we have the index weights, let's choose the portfolio weights based on dividend. You would normally calculate the weights based on trailing dividend yield, but we'll simplify this by just calculating the total dividend yield over time.
#
# Implement `calculate_dividend_weights` to return the weights for each stock based on its total dividend yield over time. This is similar to generating the weight for the index, but it's using dividend data instead.
# For example, assume the following is `dividends` data:
# ```
#                  Prices
#                A         B
# 2013-07-08     0         0
# 2013-07-09     0         1
# 2013-07-10     0.5       0
# 2013-07-11     0         0
# 2013-07-12     2         0
# ...            ...       ...
# ```
# The weights created from the function `calculate_dividend_weights` should be the following:
# ```
#                A         B
# 2013-07-08     NaN       NaN
# 2013-07-09     0         1
# 2013-07-10     0.333..   0.666..
# 2013-07-11     0.333..   0.666..
# 2013-07-12     0.714..   0.285..
# ...            ...       ...
# ```

# In[7]:

def calculate_dividend_weights(dividends):

    """
    Calculate dividend weights.

    Parameters
    ----------
    dividends : DataFrame
        Dividend for each stock and date

    Returns
    -------
    dividend_weights : DataFrame
        Weights for each stock and date
    """
    #TODO: Implement function
    divs=dividends.copy()

    weights_r=divs.cumsum()

    weights=weights_r.apply((lambda x: x/sum(x)),axis=1)
    return weights

project_tests.test_calculate_dividend_weights(calculate_dividend_weights)

# ### View Data
# Just like the index weights, let's generate the ETF weights and view them using a heatmap.

# In[8]:

etf_weights = calculate_dividend_weights(dividends)
project_helper.plot_weights(etf_weights, 'ETF Weights')

# ## Returns
# Implement `generate_returns` to generate returns data for all the stocks and dates from price data. You might notice we're implementing returns and not log returns. Since we're not dealing with volatility, we don't have to use log returns.

# In[9]:

def generate_returns(prices):
    """
    Generate returns for ticker and date.

    Parameters
    ----------
    prices : DataFrame
        Price for each ticker and date

    Returns
    -------
    returns : Dataframe
        The returns for each ticker and date
    """
    r=prices.copy()

    d=r.shift(1)

    returns=(r-d)/d

    #TODO: Implement function

    return returns

project_tests.test_generate_returns(generate_returns)

# ### View Data
# Let's generate the closing returns using `generate_returns` and view them using a heatmap.

# In[10]:

returns = generate_returns(close)
project_helper.plot_returns(returns, 'Close Returns')

# ## Weighted Returns
# With the returns of each stock computed, we can use it to compute the returns for an index or ETF. Implement `generate_weighted_returns` to create weighted returns using the returns and weights.

# In[11]:

def generate_weighted_returns(returns, weights):
    """
    Generate weighted returns.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date
    weights : DataFrame
        Weights for each ticker and date

    Returns
    -------
    weighted_returns : DataFrame
        Weighted returns for each ticker and date
    """
    weighted_returns=returns*weights

    assert returns.index.equals(weights.index)
    assert returns.columns.equals(weights.columns)

    #TODO: Implement function

    return weighted_returns

project_tests.test_generate_weighted_returns(generate_weighted_returns)

# ### View Data
# Let's generate the ETF and index returns using `generate_weighted_returns` and view them using a heatmap.

# In[12]:

index_weighted_returns = generate_weighted_returns(returns, index_weights)
etf_weighted_returns = generate_weighted_returns(returns, etf_weights)
project_helper.plot_returns(index_weighted_returns, 'Index Returns')
project_helper.plot_returns(etf_weighted_returns, 'ETF Returns')

# ## Cumulative Returns
# To compare performance between the ETF and Index, we're going to calculate the tracking error. Before we do that, we first need to calculate the index and ETF comulative returns. Implement `calculate_cumulative_returns` to calculate the cumulative returns over time given the returns.

# In[13]:

def calculate_cumulative_returns(returns):
    """
    Calculate cumulative returns.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date

    Returns
    -------
    cumulative_returns : Pandas Series
        Cumulative returns for each date
    """
    r=returns.sum(axis=1)
    ret=r.apply(lambda x: x+1)
    cum_ret=ret.cumprod()
    #TODO: Implement function

    return cum_ret

project_tests.test_calculate_cumulative_returns(calculate_cumulative_returns)

# ### View Data
# Let's generate the ETF and index cumulative returns using `calculate_cumulative_returns` and compare the two.

# In[14]:

index_weighted_cumulative_returns = calculate_cumulative_returns(index_weighted_returns)
etf_weighted_cumulative_returns = calculate_cumulative_returns(etf_weighted_returns)
project_helper.plot_benchmark_returns(index_weighted_cumulative_returns, etf_weighted_cumulative_returns, 'Smart Beta ETF vs Index')

# ## Tracking Error
# In order to check the performance of the smart beta portfolio, we can calculate the annualized tracking error against the index. Implement `tracking_error` to return the tracking error between the ETF and benchmark.
#
# For reference, we'll be using the following annualized tracking error function:
# $$ TE = \sqrt{252} * SampleStdev(r_p - r_b) $$
#
# Where $ r_p $ is the portfolio/ETF returns and $ r_b $ is the benchmark returns.
#
# _Note: When calculating the sample standard deviation, the delta degrees of freedom is 1, which is the also the default value._

# In[15]:

def tracking_error(benchmark_returns_by_date, etf_returns_by_date):
    """
    Calculate the tracking error.

    Parameters
    ----------
    benchmark_returns_by_date : Pandas Series
        The benchmark returns for each date
    etf_returns_by_date : Pandas Series
        The ETF returns for each date

    Returns
    -------
    tracking_error : float
        The tracking error
    """
    series_ret=(etf_returns_by_date)-(benchmark_returns_by_date)
    stdev=series_ret.std()
    Tracking_error=np.sqrt(252)*stdev
    assert benchmark_returns_by_date.index.equals(etf_returns_by_date.index)

    #TODO: Implement function

    return Tracking_error

project_tests.test_tracking_error(tracking_error)

# ### View Data
# Let's generate the tracking error using `tracking_error`.

# In[16]:

smart_beta_tracking_error = tracking_error(np.sum(index_weighted_returns, 1), np.sum(etf_weighted_returns, 1))
print('Smart Beta Tracking Error: {}'.format(smart_beta_tracking_error))

# # Part 2: Portfolio Optimization
#
# Now, let's create a second portfolio.  We'll still reuse the market cap weighted index, but this will be independent of the dividend-weighted portfolio that we created in part 1.
#
# We want to both minimize the portfolio variance and also want to closely track a market cap weighted index.  In other words, we're trying to minimize the distance between the weights of our portfolio and the weights of the index.
#
# $Minimize \left [ \sigma^2_p + \lambda \sqrt{\sum_{1}^{m}(weight_i - indexWeight_i)^2} \right  ]$ where $m$ is the number of stocks in the portfolio, and $\lambda$ is a scaling factor that you can choose.
#
# Why are we doing this? One way that investors evaluate a fund is by how well it tracks its index. The fund is still expected to deviate from the index within a certain range in order to improve fund performance.  A way for a fund to track the performance of its benchmark is by keeping its asset weights similar to the weights of the index.  We’d expect that if the fund has the same stocks as the benchmark, and also the same weights for each stock as the benchmark, the fund would yield about the same returns as the benchmark. By minimizing a linear combination of both the portfolio risk and distance between portfolio and benchmark weights, we attempt to balance the desire to minimize portfolio variance with the goal of tracking the index.
#
#
# ## Covariance
# Implement `get_covariance_returns` to calculate the covariance of the `returns`. We'll use this to calculate the portfolio variance.
#
# If we have $m$ stock series, the covariance matrix is an $m \times m$ matrix containing the covariance between each pair of stocks.  We can use [`Numpy.cov`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.cov.html) to get the covariance.  We give it a 2D array in which each row is a stock series, and each column is an observation at the same period of time. For any `NaN` values, you can replace them with zeros using the [`DataFrame.fillna`](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.fillna.html) function.
#
# The covariance matrix $\mathbf{P} =
# \begin{bmatrix}
# \sigma^2_{1,1} & ... & \sigma^2_{1,m} \\
# ... & ... & ...\\
# \sigma_{m,1} & ... & \sigma^2_{m,m}  \\
# \end{bmatrix}$

# In[17]:

def get_covariance_returns(returns):
    """
    Calculate covariance matrices.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date

    Returns
    -------
    returns_covariance  : 2 dimensional Ndarray
        The covariance of the returns
    """
    #TODO: Implement function

    cov=returns.fillna(0)

    return np.cov(cov,rowvar=False)

project_tests.test_get_covariance_returns(get_covariance_returns)

# ### View Data
# Let's look at the covariance generated from `get_covariance_returns`.

# In[18]:

covariance_returns = get_covariance_returns(returns)
covariance_returns = pd.DataFrame(covariance_returns, returns.columns, returns.columns)

covariance_returns_correlation = np.linalg.inv(np.diag(np.sqrt(np.diag(covariance_returns))))
covariance_returns_correlation = pd.DataFrame(
    covariance_returns_correlation.dot(covariance_returns).dot(covariance_returns_correlation),
    covariance_returns.index,
    covariance_returns.columns)

project_helper.plot_covariance_returns_correlation(
    covariance_returns_correlation,
    'Covariance Returns Correlation Matrix')

# ### portfolio variance
# We can write the portfolio variance $\sigma^2_p = \mathbf{x^T} \mathbf{P} \mathbf{x}$
#
# Recall that the $\mathbf{x^T} \mathbf{P} \mathbf{x}$ is called the quadratic form.
# We can use the cvxpy function `quad_form(x,P)` to get the quadratic form.
#
# ### Distance from index weights
# We want portfolio weights that track the index closely.  So we want to minimize the distance between them.
# Recall from the Pythagorean theorem that you can get the distance between two points in an x,y plane by adding the square of the x and y distances and taking the square root.  Extending this to any number of dimensions is called the L2 norm.  So: $\sqrt{\sum_{1}^{n}(weight_i - indexWeight_i)^2}$  Can also be written as $\left \| \mathbf{x} - \mathbf{index} \right \|_2$.  There's a cvxpy function called [norm()](https://www.cvxpy.org/api_reference/cvxpy.atoms.other_atoms.html#norm)
# `norm(x, p=2, axis=None)`.  The default is already set to find an L2 norm, so you would pass in one argument, which is the difference between your portfolio weights and the index weights.
#
# ### objective function
# We want to minimize both the portfolio variance and the distance of the portfolio weights from the index weights.
# We also want to choose a `scale` constant, which is $\lambda$ in the expression.
#
# $\mathbf{x^T} \mathbf{P} \mathbf{x} + \lambda \left \| \mathbf{x} - \mathbf{index} \right \|_2$
#
#
# This lets us choose how much priority we give to minimizing the difference from the index, relative to minimizing the variance of the portfolio.  If you choose a higher value for `scale` ($\lambda$).
#
# We can find the objective function using cvxpy `objective = cvx.Minimize()`.  Can you guess what to pass into this function?
#
#

# ### constraints
# We can also define our constraints in a list.  For example, you'd want the weights to sum to one. So $\sum_{1}^{n}x = 1$.  You may also need to go long only, which means no shorting, so no negative weights.  So $x_i >0 $ for all $i$. you could save a variable as `[x >= 0, sum(x) == 1]`, where x was created using `cvx.Variable()`.
#
# ### optimization
# So now that we have our objective function and constraints, we can solve for the values of $\mathbf{x}$.
# cvxpy has the constructor `Problem(objective, constraints)`, which returns a `Problem` object.
#
# The `Problem` object has a function solve(), which returns the minimum of the solution.  In this case, this is the minimum variance of the portfolio.
#
# It also updates the vector $\mathbf{x}$.
#
# We can check out the values of $x_A$ and $x_B$ that gave the minimum portfolio variance by using `x.value`

# In[19]:

import cvxpy as cvx

def get_optimal_weights(covariance_returns, index_weights, scale=2.0):
    """
    Find the optimal weights.

    Parameters
    ----------
    covariance_returns : 2 dimensional Ndarray
        The covariance of the returns
    index_weights : Pandas Series
        Index weights for all tickers at a period in time
    scale : int
        The penalty factor for weights the deviate from the index
    Returns
    -------
    x : 1 dimensional Ndarray
        The solution for x
    """

    m=index_weights.shape[0]

    x=cvx.Variable(m)

    p = covariance_returns

    portfolio_variance = cvx.quad_form(x,p)

    diff=(x-index_weights)

    dist=cvx.norm(diff,p=2,axis=None)

    objective = cvx.Minimize(portfolio_variance + scale*(dist))

    constraints = [x >= 0 , sum(x) == 1]

    problem= cvx.Problem(objective,constraints)

    solved= problem.solve()

    x_vals= x.value

    assert len(covariance_returns.shape) == 2
    assert len(index_weights.shape) == 1
    assert covariance_returns.shape[0] == covariance_returns.shape[1]  == index_weights.shape[0]

    #TODO: Implement function

    return x_vals

project_tests.test_get_optimal_weights(get_optimal_weights)

# ## Optimized Portfolio
# Using the `get_optimal_weights` function, let's generate the optimal ETF weights without rebalanceing. We can do this by feeding in the covariance of the entire history of data. We also need to feed in a set of index weights. We'll go with the average weights of the index over time.

# In[20]:

raw_optimal_single_rebalance_etf_weights = get_optimal_weights(covariance_returns.values, index_weights.iloc[-1])
optimal_single_rebalance_etf_weights = pd.DataFrame(
    np.tile(raw_optimal_single_rebalance_etf_weights, (len(returns.index), 1)),
    returns.index,
    returns.columns)

# With our ETF weights built, let's compare it to the index. Run the next cell to calculate the ETF returns and compare it to the index returns.

# In[21]:

optim_etf_returns = generate_weighted_returns(returns, optimal_single_rebalance_etf_weights)
optim_etf_cumulative_returns = calculate_cumulative_returns(optim_etf_returns)
project_helper.plot_benchmark_returns(index_weighted_cumulative_returns, optim_etf_cumulative_returns, 'Optimized ETF vs Index')

optim_etf_tracking_error = tracking_error(np.sum(index_weighted_returns, 1), np.sum(optim_etf_returns, 1))
print('Optimized ETF Tracking Error: {}'.format(optim_etf_tracking_error))

# ## Rebalance Portfolio Over Time
# The single optimized ETF portfolio used the same weights for the entire history. This might not be the optimal weights for the entire period. Let's rebalance the portfolio over the same period instead of using the same weights. Implement `rebalance_portfolio` to rebalance a portfolio.
#
# Reblance the portfolio every n number of days, which is given as `shift_size`. When rebalancing, you should look back a certain number of days of data in the past, denoted as `chunk_size`. Using this data, compute the optoimal weights using `get_optimal_weights` and `get_covariance_returns`.

# In[22]:

def rebalance_portfolio(returns, index_weights, shift_size, chunk_size):
    """
    Get weights for each rebalancing of the portfolio.

    Parameters
    ----------
    returns : DataFrame
        Returns for each ticker and date
    index_weights : DataFrame
        Index weight for each ticker and date
    shift_size : int
        The number of days between each rebalance
    chunk_size : int
        The number of days to look in the past for rebalancing

    Returns
    -------
    all_rebalance_weights  : list of Ndarrays
        The ETF weights for each point they are rebalanced
    """
    returns = returns.fillna(0).copy()

    generate=np.arange(1,len(returns))

    lists_i=[]

    count=0

    lists_r=[]

    count=-1
    for i in generate:
        count+=1
        if len(lists_r)!=chunk_size:
            lists_r.append(count)

        if len(lists_r)==chunk_size:
            n = lists_r[shift_size-chunk_size:]
            count=n[-1]
            lists_i.append(lists_r)
            lists_r = []
            for i in n:
                lists_r.append(i)
    lists_d=[]

    for d in lists_i :
        lists_d.append(d[-1])

    list_vals=[]
    for i in lists_i:
        list_input=[]
        for n in i:
            list_input.append(returns.iloc[n].values)

        list_vals.append(list_input)
        list_input=[]

    v=np.array(list_vals)

    rest = []
    for i in v:
        k= pd.DataFrame(i)
        rest.append(get_covariance_returns(k))

    weights = []

    for (i,n) in zip(lists_d,rest):
        index_w = index_weights.iloc[i]
        weights.append(get_optimal_weights(n,index_w,scale=2.0))

    assert returns.index.equals(index_weights.index)
    assert returns.columns.equals(index_weights.columns)
    assert shift_size > 0
    assert chunk_size >= 0

    #TODO: Implement function

    return weights

project_tests.test_rebalance_portfolio(rebalance_portfolio)

# Run the following cell to rebalance the portfolio using `rebalance_portfolio`.

# In[23]:

chunk_size = 250
shift_size = 5
all_rebalance_weights = rebalance_portfolio(returns, index_weights, shift_size, chunk_size)

# ## Portfolio Turnover
# With the portfolio rebalanced, we need to use a metric to measure the cost of rebalancing the portfolio. Implement `get_portfolio_turnover` to calculate the annual portfolio turnover. We'll be using the formulas used in the classroom:
#
# $ AnnualizedTurnover =\frac{SumTotalTurnover}{NumberOfRebalanceEvents} * NumberofRebalanceEventsPerYear $
#
# $ SumTotalTurnover =\sum_{t,n}{\left | x_{t,n} - x_{t+1,n} \right |} $ Where $ x_{t,n} $ are the weights at time $ t $ for equity $ n $.
#
# $ SumTotalTurnover $ is just a different way of writing $ \sum \left | x_{t_1,n} - x_{t_2,n} \right | $

# In[40]:

def get_portfolio_turnover(all_rebalance_weights, shift_size, rebalance_count, n_trading_days_in_year=252):
    """
    Calculage portfolio turnover.

    Parameters
    ----------
    all_rebalance_weights : list of Ndarrays
        The ETF weights for each point they are rebalanced
    shift_size : int
        The number of days between each rebalance
    rebalance_count : int
        Number of times the portfolio was rebalanced
    n_trading_days_in_year: int
        Number of trading days in a year

    Returns
    -------
    portfolio_turnover  : float
        The portfolio turnover
    """
    rebalance_e = n_trading_days_in_year/shift_size

    df = pd.DataFrame(all_rebalance_weights)
    df_s=df.diff(periods=-1)
    df_t=df_s.fillna(0)

    turnover_per_n = df_t.apply( lambda x: abs(x)).sum()

    sum_total_turnover = turnover_per_n.sum()

    annualized_turnover = (sum_total_turnover*rebalance_e)/rebalance_count

    assert shift_size > 0
    assert rebalance_count > 0

    #TODO: Implement function

    return annualized_turnover

project_tests.test_get_portfolio_turnover(get_portfolio_turnover)

# Run the following cell to get the portfolio turnover from  `get_portfolio turnover`.

# In[41]:

print(get_portfolio_turnover(all_rebalance_weights, shift_size, len(all_rebalance_weights) - 1))

# That's it! You've built a smart beta portfolio in part 1 and did portfolio optimization in part 2. You can now submit your project.

# ## Submission
# Now that you're done with the project, it's time to submit it. Click the submit button in the bottom right. One of our reviewers will give you feedback on your project with a pass or not passed grade. You can continue to the next section while you wait for feedback.


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

# # K-mer
# Basic K-mer counting.

# In[1]:

import time
def show_time():
    t = time.time()
    print(time.strftime('%Y-%m-%d %H:%M:%S %Z', time.localtime(t)))
show_time()

# In[2]:

PC_SEQUENCES=32000
NC_SEQUENCES=32000
RNA_LEN=32
CDS_LEN=16

# In[3]:

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

# In[4]:

import sys
IN_COLAB = False
try:
    from google.colab import drive
    IN_COLAB = True
except:
    pass
if IN_COLAB:
    print("On Google CoLab, mount cloud-local file, get our code from GitHub.")
    PATH='/content/drive/'
    #drive.mount(PATH,force_remount=True)  # hardly ever need this
    #drive.mount(PATH)    # Google will require login credentials
    DATAPATH=PATH+'My Drive/data/'  # must end in "/"
    import requests
    r = requests.get('https://raw.githubusercontent.com/ShepherdCode/Soars2021/master/SimTools/RNA_describe.py')
    with open('RNA_describe.py', 'w') as f:
        f.write(r.text)
    from RNA_describe import Random_Base_Oracle
else:
        print("CoLab not working. On my PC, use relative paths.")
        DATAPATH='data/'  # must end in "/"
        sys.path.append("..") # append parent dir in order to use sibling dirs
        from SimTools.RNA_describe import Random_Base_Oracle
MODELPATH="BestModel"  # saved on cloud instance and lost after logout
#MODELPATH=DATAPATH+MODELPATH  # saved on Google Drive but requires login

# ## K-mer counting

# ### Functions to create the dict of {kmer:count}

# In[5]:

def make_kmer_keys(K):
    shorter_kmers=['']
    for i in range(K):
        longer_kmers=[]
        for mer in shorter_kmers:
            # No support for N or any non-ACGT bases.
            longer_kmers.append(mer+'A')
            longer_kmers.append(mer+'C')
            longer_kmers.append(mer+'G')
            longer_kmers.append(mer+'T')
        shorter_kmers = longer_kmers
    return shorter_kmers
def make_kmer_dict(keys,init=0):
    return dict.fromkeys(keys,init)
def make_dict_upto_K(max_K):
    keys=make_kmer_keys(1)
    for k in range(2,max_K+1):
        keys.extend(make_kmer_keys(k))
    counts = make_kmer_dict(keys)
    return counts

# ### Naive K-mer counting algorithm
# Algorithm:
# 1. for every string
#     1. for every K
#         1. for every position
#             1. kmer=substring
#             2. count{kmer}++

# In[6]:

def update_count_one_K(counts,K,rna,tail=False):
    L = len(rna)
    padding=" "*(K-1)
    padded=rna+padding
    for i in range(0,L-K+1):
        kmer=padded[i:i+K]
        counts[kmer] += 1
    if tail and K>1:
        # for Harvester algorithm, count last letters as special case
        for start_pos in range(L-K+1,L):
            for end_pos in range(start_pos+1,L+1):
                kmer=rna[start_pos:end_pos]
                counts[kmer] += 1
    return counts
def update_count_upto_K(counts,max_K,sample,tail=False):
    for i in range(1,max_K+1):
        update_count_one_K(counts,i,sample,tail)
    return counts

# ### Harvester K-mer counting algorithm
# Algorithm:
# 1. Count K-mers for max K only
# 2. For each K-mer in counts table:
#     1. For every prefix of the K-mer:
#         1. count{prefix} += count{kmer}
# 3. Handle last K-1 letters of each string as special case

# In[7]:

def harvest_counts_from_K(counts,max_K):
    for kmer in counts.keys():
        klen = len(kmer)
        kcnt = counts[kmer]
        if klen==max_K and kcnt>0:
            for i in range(1,klen):
                prefix = kmer[:i]
                counts[prefix] += kcnt
    return counts

# In[8]:

def count_to_frequency(counts,max_K):
    freqs = dict.fromkeys(counts.keys(),0.0)
    for k in range(1,max_K+1):
        tot = 0
        for kmer in counts.keys():
            if len(kmer)==k:
                tot += counts[kmer]
        for kmer in counts.keys():
            if len(kmer)==k:
                freqs[kmer] = 1.0*counts[kmer]/tot
    return freqs

# ## Demo

# ### Demo: Naive algorithm

# In[9]:

MAX_K = 3
counts1 = make_dict_upto_K(MAX_K)
print("Initial counts:\n",counts1)

sample = "ACCGGGTTTTACGTACGT"
update_count_upto_K(counts1,MAX_K,sample)
print("Final counts:\n",counts1)

# ### Demo: Harvester algorithm

# In[10]:

MAX_K = 3
counts2 = make_dict_upto_K(MAX_K)
print("Initial counts:\n",counts2)

sample = "ACCGGGTTTTACGTACGT"
update_count_one_K(counts2,MAX_K,sample,True)
print("Partial counts (just max K and special case letters)\n:",counts2)
harvest_counts_from_K(counts2,MAX_K)
print("Final counts (includes smaller values of K):\n",counts2)

# In[11]:

if counts1==counts2:
    print("Success. Harvester output matches naive results!")
else:
    print("Fail. Harvester output differs from naive results!")

# In[12]:

freqs = count_to_frequency(counts2,MAX_K)
print ("Frequency:\n",freqs)

# ## Demo on large dataset

# In[13]:

rbo=Random_Base_Oracle(RNA_LEN,True)
pc_all,nc_all = rbo.get_partitioned_sequences(CDS_LEN,10) # just testing
pc_all,nc_all = rbo.get_partitioned_sequences(CDS_LEN,PC_SEQUENCES)
print("Use",len(pc_all),"PC seqs")
print("Use",len(nc_all),"NC seqs")

# In[14]:

MAX_K = 3
pc_counts = make_dict_upto_K(MAX_K)
for sample in pc_all:
    update_count_one_K(pc_counts,MAX_K,sample,True)
harvest_counts_from_K(pc_counts,MAX_K)
print("PC counts:\n",pc_counts)
pc_freqs = count_to_frequency(pc_counts,MAX_K)
print ("Frequency:\n",pc_freqs)

# In[15]:

nc_counts = make_dict_upto_K(MAX_K)
for sample in nc_all:
    update_count_one_K(nc_counts,MAX_K,sample,True)
harvest_counts_from_K(nc_counts,MAX_K)
print("NC counts:\n",nc_counts)
nc_freqs = count_to_frequency(nc_counts,MAX_K)
print ("Frequency:\n",nc_freqs)

# In[ ]:


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

# ## Rover Lab Notebook
# This notebook contains the functions from the lesson and provides the scaffolding you need to test out your mapping methods.  The steps you need to complete in this notebook for the project are the following:
#
# * First just run each of the cells in the notebook, examine the code and the results of each.
#
# **Note: For the online lab, data has been collected and provided for you. If you would like to try locally please do so! Please continue instructions from the continue point.**
# * Run the simulator in "Training Mode" and record some data. Note: the simulator may crash if you try to record a large (longer than a few minutes) dataset, but you don't need a ton of data, just some example images to work with.
# * Change the data directory path (2 cells below) to be the directory where you saved data
# * Test out the functions provided on your data
#
# **Continue Point**
# * Write new functions (or modify existing ones) to report and map out detections of obstacles and rock samples (yellow rocks)
# * Populate the `process_image()` function with the appropriate steps/functions to go from a raw image to a worldmap.
# * Run the cell that calls `process_image()` using `moviepy` functions to create video output
# * Once you have mapping working, move on to modifying `perception.py` and `decision.py` in the project to allow your rover to navigate and map in autonomous mode!
#
# **Note: If, at any point, you encounter frozen display windows or other confounding issues, you can always start again with a clean slate by going to the "Kernel" menu above and selecting "Restart & Clear Output".**
#
# **Run the next cell to get code highlighting in the markdown cells.**

# In[31]:

get_ipython().run_cell_magic('HTML', '', '<style> code {background-color : orange !important;} </style>')

# In[32]:

get_ipython().run_line_magic('matplotlib', 'inline')
#%matplotlib qt # Choose %matplotlib qt to plot to an interactive window (note it may show up behind your browser)
# Make some of the relevant imports
import cv2 # OpenCV for perspective transform
import numpy as np
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import scipy.misc # For saving images as needed
import glob  # For reading in a list of images from a folder

# ## Quick Look at the Data
# There's some example data provided in the `test_dataset` folder.  This basic dataset is enough to get you up and running but if you want to hone your methods more carefully you should record some data of your own to sample various scenarios in the simulator.
#
# Next, read in and display a random image from the `test_dataset` folder

# In[33]:

path = './test_dataset/IMG/*'
img_list = glob.glob(path)
# Grab a random image and display it
idx = np.random.randint(0, len(img_list)-1)
image = mpimg.imread(img_list[idx])
plt.imshow(image)

# ## Calibration Data
# Read in and display example grid and rock sample calibration images.  You'll use the grid for perspective transform and the rock image for creating a new color selection that identifies these samples of interest.

# In[34]:

# In the simulator you can toggle on a grid on the ground for calibration
# You can also toggle on the rock samples with the 0 (zero) key.
# Here's an example of the grid and one of the rocks
example_grid = './calibration_images/example_grid1.jpg'
example_rock = './calibration_images/example_rock1.jpg'
grid_img = mpimg.imread(example_grid)
rock_img = mpimg.imread(example_rock)

fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(grid_img)
plt.subplot(122)
plt.imshow(rock_img)

# ## Perspective Transform
#
# Define the perspective transform function from the lesson and test it on an image.

# In[35]:

# Define a function to perform a perspective transform
# I've used the example grid image above to choose source points for the
# grid cell in front of the rover (each grid cell is 1 square meter in the sim)
# Define a function to perform a perspective transform
def perspect_transform(img, src, dst):
    M = cv2.getPerspectiveTransform(src, dst)
    warped = cv2.warpPerspective(img, M, (img.shape[1], img.shape[0]))# keep same size as input image
    mask = cv2.warpPerspective(np.ones_like(img[:,:,0]), M, (img.shape[1], img.shape[0]))
    return warped, mask

# Define calibration box in source (actual) and destination (desired) coordinates
# These source and destination points are defined to warp the image
# to a grid where each 10x10 pixel square represents 1 square meter
# The destination box will be 2*dst_size on each side
dst_size = 5
# Set a bottom offset to account for the fact that the bottom of the image
# is not the position of the rover but a bit in front of it
# this is just a rough guess, feel free to change it!
bottom_offset = 6
source = np.float32([[14, 140], [301 ,140],[200, 96], [118, 96]])
destination = np.float32([[image.shape[1]/2 - dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  [image.shape[1]/2 - dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  ])
warped, mask = perspect_transform(image, source, destination)
fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(warped)
plt.subplot(122)
plt.imshow(mask)
#scipy.misc.imsave('../output/warped_example.jpg', warped)

# ## Color Thresholding
# Define the color thresholding function from the lesson and apply it to the warped image
#
# **TODO:** Ultimately, you want your map to not just include navigable terrain but also obstacles and the positions of the rock samples you're searching for.  Modify this function or write a new function that returns the pixel locations of obstacles (areas below the threshold) and rock samples (yellow rocks in calibration images), such that you can map these areas into world coordinates as well.
# **Suggestion:** Think about imposing a lower and upper boundary in your color selection to be more specific about choosing colors.  Feel free to get creative and even bring in functions from other libraries.  Here's an example of [color selection](http://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_colorspaces/py_colorspaces.html) using OpenCV.
# **Beware:** if you start manipulating images with OpenCV, keep in mind that it defaults to `BGR` instead of `RGB` color space when reading/writing images, so things can get confusing.

# In[36]:

# Identify pixels above the threshold
# Threshold of RGB > 160 does a nice job of identifying ground pixels only
def color_thresh(img, rgb_thresh=(160, 160, 160)):
    # Create an array of zeros same xy size as img, but single channel
    color_select = np.zeros_like(img[:,:,0])
    # Require that each pixel be above all three threshold values in RGB
    # above_thresh will now contain a boolean array with "True"
    # where threshold was met
    above_thresh = (img[:,:,0] > rgb_thresh[0])                 & (img[:,:,1] > rgb_thresh[1])                 & (img[:,:,2] > rgb_thresh[2])
    # Index the array of zeros with the boolean array and set to 1
    color_select[above_thresh] = 1

    # Return the binary image
    return color_select

def yellow_rock_thresh(img):
    # identify if the image has a rock in it
    # define a threshold for the rock object
    color_select_yellow = np.zeros_like(img[:,:,0])
    yellow_thresh = (img[:,:,0] > 50) & (img[:,:,0] < 258)                  & (img[:,:,1] > 115) & (img[:,:,1] < 256)                  & (img[:,:,2] > -1) &  (img[:,:,2] < 100)

    color_select_yellow[yellow_thresh] = 1

    is_yellow_rock = False
    for nested_list in color_select_yellow:
        for item in nested_list:
            if item:
                is_yellow_rock = True
                break
    return color_select_yellow, is_yellow_rock

threshed = color_thresh(warped)
yellow_rock_img, is_yellow_rock = yellow_rock_thresh(warped)
fig = plt.figure(figsize=(12,3))
plt.subplot(121)
plt.imshow(threshed, cmap='gray')
plt.subplot(122)
plt.imshow(yellow_rock_img, cmap='gray')
print (is_yellow_rock)
#scipy.misc.imsave('../output/warped_threshed.jpg', threshed*255)

# ## Coordinate Transformations
# Define the functions used to do coordinate transforms and apply them to an image.

# In[37]:

def rover_coords(binary_img):
    # Identify nonzero pixels
    ypos, xpos = binary_img.nonzero()
    # Calculate pixel positions with reference to the rover position being at the
    # center bottom of the image.
    x_pixel = np.absolute(ypos - binary_img.shape[0]).astype(np.float)
    y_pixel = -(xpos - binary_img.shape[0]).astype(np.float)
    return x_pixel, y_pixel

# Define a function to convert to radial coords in rover space
def to_polar_coords(x_pixel, y_pixel):
    # Convert (x_pixel, y_pixel) to (distance, angle)
    # in polar coordinates in rover space
    # Calculate distance to each pixel
    dist = np.sqrt(x_pixel**2 + y_pixel**2)
    # Calculate angle away from vertical for each pixel
    angles = np.arctan2(y_pixel, x_pixel)
    return dist, angles

# Define a function to apply a rotation to pixel positions
def rotate_pix(xpix, ypix, yaw):
    # TODO:
    # Convert yaw to radians
    yaw_rad = yaw * np.pi / 180

    # Apply a rotation
    x_rotated = xpix * np.cos(yaw_rad) - ypix * np.sin(yaw_rad)
    y_rotated = xpix * np.sin(yaw_rad) + ypix * np.cos(yaw_rad)
    # Return the result
    return x_rotated, y_rotated

# Define a function to perform a translation
def translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale):
    # TODO:
    # Apply a scaling and a translation
    # Assume a scale factor of 10 between world space pixels and rover space pixels
    # Perform translation and convert to integer since pixel values can't be float
    x_world = np.int_(xpos + (xpix_rot / scale))
    y_world = np.int_(ypos + (ypix_rot / scale))

    # Return the result
    return x_world, y_world

# Define a function to apply rotation and translation (and clipping)
# Once you define the two functions above this function should work
def pix_to_world(xpix, ypix, xpos, ypos, yaw, world_size, scale):
    # Apply rotation
    xpix_rot, ypix_rot = rotate_pix(xpix, ypix, yaw)
    # Apply translation
    xpix_tran, ypix_tran = translate_pix(xpix_rot, ypix_rot, xpos, ypos, scale)
    # Perform rotation, translation and clipping all at once
    x_pix_world = np.clip(np.int_(xpix_tran), 0, world_size - 1)
    y_pix_world = np.clip(np.int_(ypix_tran), 0, world_size - 1)
    # Return the result
    return x_pix_world, y_pix_world

# Grab another random image
idx = np.random.randint(0, len(img_list)-1)
#image = mpimg.imread(img_list[idx])
image = rock_img
warped, mask = perspect_transform(image, source, destination)
threshed = color_thresh(warped)  # navigable terrain pixels
obs_map = np.absolute(np.float32(threshed)-1)*mask  #obstacles pixels

# Calculate pixel values in rover-centric coords and distance/angle to all pixels
xpix, ypix = rover_coords(threshed)
dist, angles = to_polar_coords(xpix, ypix)
mean_dir = np.mean(angles)

# Do some plotting
fig = plt.figure(figsize=(12,9))
plt.subplot(221)
plt.imshow(image)
plt.subplot(222)
plt.imshow(warped)
plt.subplot(223)
plt.imshow(threshed, cmap='gray')
plt.subplot(224)
plt.plot(xpix, ypix, '.')
plt.ylim(-160, 160)
plt.xlim(0, 160)
arrow_length = 100
x_arrow = arrow_length * np.cos(mean_dir)
y_arrow = arrow_length * np.sin(mean_dir)
plt.arrow(0, 0, x_arrow, y_arrow, color='red', zorder=2, head_width=10, width=2)

# ## Read in saved data and ground truth map of the world
# The next cell is all setup to read your saved data into a `pandas` dataframe.  Here you'll also read in a "ground truth" map of the world, where white pixels (pixel value = 1) represent navigable terrain.
#
# After that, we'll define a class to store telemetry data and pathnames to images.  When you instantiate this class (`data = Databucket()`) you'll have a global variable called `data` that you can refer to for telemetry and map data within the `process_image()` function in the following cell.
#

# In[38]:

# Import pandas and read in csv file as a dataframe
import pandas as pd
# Change this path to your data directory
df = pd.read_csv('./test_dataset/robot_log.csv')
csv_img_list = df["Path"].tolist() # Create list of image pathnames
# Read in ground truth map and create a 3-channel image with it
ground_truth = mpimg.imread('./calibration_images/map_bw.png')
ground_truth_3d = np.dstack((ground_truth*0, ground_truth*255, ground_truth*0)).astype(np.float)

# Creating a class to be the data container
# Will read in saved data from csv file and populate this object
# Worldmap is instantiated as 200 x 200 grids corresponding
# to a 200m x 200m space (same size as the ground truth map: 200 x 200 pixels)
# This encompasses the full range of output position values in x and y from the sim
class Databucket():
    def __init__(self):
        self.images = csv_img_list
        self.xpos = df["X_Position"].values
        self.ypos = df["Y_Position"].values
        self.yaw = df["Yaw"].values
        self.count = -1 # This will be a running index, setting to -1 is a hack
                        # because moviepy (below) seems to run one extra iteration
        self.worldmap = np.zeros((200, 200, 3)).astype(np.float)
        self.ground_truth = ground_truth_3d # Ground truth worldmap

# Instantiate a Databucket().. this will be a global variable/object
# that you can refer to in the process_image() function below
data = Databucket()

# ## Write a function to process stored images
#
# Modify the `process_image()` function below by adding in the perception step processes (functions defined above) to perform image analysis and mapping.  The following cell is all set up to use this `process_image()` function in conjunction with the `moviepy` video processing package to create a video from the images you saved taking data in the simulator.
#
# In short, you will be passing individual images into `process_image()` and building up an image called `output_image` that will be stored as one frame of video.  You can make a mosaic of the various steps of your analysis process and add text as you like (example provided below).
#
#
#
# To start with, you can simply run the next three cells to see what happens, but then go ahead and modify them such that the output video demonstrates your mapping process.  Feel free to get creative!

# In[43]:

# Define a function to pass stored images to
# reading rover position and yaw angle from csv file
# This function will be used by moviepy to create an output video
def process_image(img):
    # Example of how to use the Databucket() object defined above
    # to print the current x, y and yaw values
    # print(data.xpos[data.count], data.ypos[data.count], data.yaw[data.count])

    # TODO:
    # 1) Define source and destination points for perspective transform
    source = np.float32([
                [14, 140], [301, 140], # bottom left to right
                [200, 96], [118, 96]]) # top rigt to left
    dst_size = 5
    bottom_offset = 6
    destination = np.float32([[image.shape[1]/2 - dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - bottom_offset],
                  [image.shape[1]/2 + dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  [image.shape[1]/2 - dst_size, image.shape[0] - 2*dst_size - bottom_offset],
                  ])

    # 2) Apply perspective transform
    warped, mask = perspect_transform(img, source, destination)
    # 3) Apply color threshold to identify navigable terrain/obstacles/rock samples

    # def color_thresh
    threshed = color_thresh(warped) # navigable pixels
    yellow_rock_threshed, is_yellow_rock = yellow_rock_thresh(warped)  # rock pixels
    obs_map = np.absolute(np.float32(threshed)-1)*mask  #obstacles pixels

    # 4) Convert thresholded image pixel values to rover-centric coords
    xpix, ypix = rover_coords(threshed)
    xpix_obs, ypix_obs = rover_coords(obs_map)

    # 5) Convert rover-centric pixel values to world coords
    scale = 10
    x_world, y_world = pix_to_world(xpix, ypix, data.xpos[data.count], data.ypos[data.count], \
        data.yaw[data.count], data.worldmap.shape[0], scale)
    x_obs, y_obs = pix_to_world(xpix_obs, ypix_obs, data.xpos[data.count], \
        data.ypos[data.count], data.yaw[data.count], data.worldmap.shape[0], scale)

    # 6) Update worldmap (to be displayed on right side of screen)
        # Example: data.worldmap[obstacle_y_world, obstacle_x_world, 0] += 1
        #          data.worldmap[rock_y_world, rock_x_world, 1] += 1
        #          data.worldmap[navigable_y_world, navigable_x_world, 2] += 1
    data.worldmap[y_world, x_world, 2] += 255  # blue nav terrain
    data.worldmap[y_obs, x_obs, 0] += 255 # red obstacles
    nav_pix = data.worldmap[:,:,2] > 0
    data.worldmap[nav_pix, 0] = 0

    if is_yellow_rock:
        xpix_rock, ypix_rock = rover_coords(yellow_rock_threshed)
        x_world_rock, y_world_rock = pix_to_world(xpix_rock, ypix_rock, data.xpos[data.count], \
            data.ypos[data.count], data.yaw[data.count], data.worldmap.shape[0], scale)
        data.worldmap[y_world_rock, x_world_rock, :] = 255

    # 7) Make a mosaic image, below is some example code
        # First create a blank image (can be whatever shape you like)
    output_image = np.zeros((img.shape[0] + data.worldmap.shape[0], img.shape[1]*2, 3))
        # Next you can populate regions of the image with various output
        # Here I'm putting the original image in the upper left hand corner
    #left corner
    output_image[0:img.shape[0], 0:img.shape[1]] = img

    # first a warped image
    # upper right
    output_image[0:img.shape[0], img.shape[1]:] = warped

    # Overlay worldmap with ground truth map
    map_add = cv2.addWeighted(data.worldmap, 1, data.ground_truth, 0.5, 0)
        # Flip map overlay so y-axis points upward and add to output_image
    #lower left
    output_image[img.shape[0]:, 0:data.worldmap.shape[1]] = np.flipud(map_add)
    #lower right ??????

        # Then putting some text over the image
    if is_yellow_rock:
        cv2.putText(output_image,"Rock Found!!!", (20, 20),
                cv2.FONT_HERSHEY_COMPLEX, 0.4, (255, 255, 255), 1)
    else:
         cv2.putText(output_image,"No Rock", (20, 20),
                cv2.FONT_HERSHEY_COMPLEX, 0.4, (255, 255, 255), 1)
    data.count += 1 # Keep track of the index in the Databucket()

    return output_image

# ## Make a video from processed image data
# Use the [moviepy](https://zulko.github.io/moviepy/) library to process images and create a video.
#

# In[44]:

# Import everything needed to edit/save/watch video clips
from moviepy.editor import VideoFileClip
from moviepy.editor import ImageSequenceClip

# Define pathname to save the output video
output = './output/test_mapping.mp4'
data = Databucket() # Re-initialize data in case you're running this cell multiple times
clip = ImageSequenceClip(data.images, fps=60) # Note: output video will be sped up because
                                          # recording rate in simulator is fps=25
new_clip = clip.fl_image(process_image) #NOTE: this function expects color images!!
get_ipython().run_line_magic('time', 'new_clip.write_videofile(output, audio=False)')

# ### This next cell should function as an inline video player
# If this fails to render the video, try running the following cell (alternative video rendering method).  You can also simply have a look at the saved mp4 in your `/output` folder

# In[45]:

output = './output/test_mapping.mp4'
from IPython.display import HTML
HTML("""
<video width="960" height="540" controls>
  <source src="{0}">
</video>
""".format(output))

# In[ ]:

# In[ ]:


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Import Main Librairies

# +
# Basic Python modules
import random
import csv
import pickle
import os

# Scientific modules
import numpy as np
import scipy as scp
import matplotlib
from matplotlib import pyplot as plt

from scipy.io import loadmat
import networkx as nx

# For better looking graphs
import seaborn as sns
sns.set()

# Just a custom color palette that I use
# blue, darkblue, red, orange, green, palegreen, yellow, brokenwhite, brokengrey
colors = ["#AF3127","#6182B5","#112A3C","#D99C37","#90A954","#C5B868","#FAC764","#DAC0A6","#C4C2D1"]
sns.set_palette(sns.color_palette(colors))
# -

# ## Configure paths for data once for all

# +
#generate_images = False

#pickle_data_folder = 'new_new_pickle_data_cross/'

# Folder in which data is placed
data_folder = '../fake_data'

# Folder in which some computed objects will be saved
pickle_folder = 'cross_validation'

# Control whether we recompute data even if it has been previously
# computed and pickled
recompute_pickle = False

# Suffix at the end of each .mat file; it is specified here in order to select
# relevant files, as well as to make file name lighter during loading for
# further operations (such as printing subjects names), since it does not carry
# additional information.
suffix = '_fiber_number.mat'

# For instance here, with those setting, every ../fake_data/*_fiber_number.mat
# will be loaded

# Keys used to split data between patients and controls. Subject whose filename
# contains one of the control_keys will be affected to the control cohort, and
# similarly for patients.
control_keys = ['060', 'dep', 'dpr', 'S', 'TI']
patient_keys = ['lgp']

# By default, the code expects a "table.csv" present in data_folder, containing
# information about the patients, such as their age, the duration of the
# disease, etc.
csv_path = data_folder + "/table.csv"

# -

# ## Load data

# +
def get_matrix_file_list(data_folder, suffix):
    """
        Return the list of files in the folder data_folder ending by suffix
    """
    file_list = [f for f in os.listdir(data_folder) if f.endswith(suffix)]
    return list(map(lambda x: data_folder + '/' + x, file_list))

def load_matrix_file(file):
    """
        Return the matrix loaded from a Matlab data file. Note that the
        'Measure' key is hardcoded, so you might need to adjust it to your own
        data.
    """
    return loadmat(file)['Measure']

# Eventually create the folder in which data will be saved
if not os.path.exists(pickle_folder):
    os.mkdir(pickle_folder)

# Create a dictionnary of all the matrices, where each matrix gets associated to
# the filename of the corresponding .mat file minus the suffix.
connectivity_matrices = {}

for f in get_matrix_file_list(data_folder, suffix):
    connectivity_matrices[f.replace(suffix,'').replace(data_folder + '/','')] = load_matrix_file(f)

# Create a dictionnary of metadata for each patient, obtained from the file
# at csv_path (by default 'data_foler/table.csv')
patient_info_dict = {}

with open(csv_path, 'r') as csv_file:
    metadata = csv.DictReader(csv_file)
    # Each patient is associated to a dictionnary containing all its information
    for row in metadata:
        metadata_dict = {key:row[key] for key in metadata.fieldnames if key != 'Subject'}
        patient_info_dict[row['Subject']] = metadata_dict

print("Succesfully loaded {} matrices from {}.".format(len(connectivity_matrices), data_folder))
print("Metadata has been found in {} for {} subjects.".format(csv_path,len(patient_info_dict)))
# -

# ## Split Data into Cohorts

# +
# list of controls and patients names
controls = []
patients = []

# The following can be used to limit the number of either controls or patients
# considered for the study if they are set to some non infinite number.
controls_count = np.inf
patients_count = np.inf

current_control = 0
current_patient = 0

for key in [*connectivity_matrices]:
    # Use patients_keys and control_keys list to classify subject into cohorts
    if any(list(map(lambda x: x in key, patient_keys))) and current_patient < patients_count:
        patients.append(key)
        current_patient += 1
    elif any(list(map(lambda x: x in key, control_keys))) and current_control < controls_count:
        controls.append(key)
        current_control += 1
    else:
        print("Patient {} cannot be classified either as control or patient.".format(key))

controls_count = current_control
patients_count = current_patient

subject_count = len(patients) + len(controls)

print("Classified {} controls and {} patients (total {} subjects)".format(controls_count, patients_count, subject_count))

# -

# ## Basic network manipulation functions

# +
def plot_connectivity_matrix(*matrix, save=None):
    """
        Plot the connectivity matrix as a colormap
        It uses a logarithmic scale for the colors and it is possible to plot
        multiple matrices at once.
    """

    max_rows = 2
    line = len(matrix) / 3 + 1

    for i, mat in enumerate(matrix):
        plt.subplot(line, max_rows, i+1)
        plt.imshow(np.log(mat+1),cmap='viridis')

    plt.tight_layout()
    if save is None:
        plt.show()
    else:
        plt.savefig(save,bbox_inches='tight')

def get_network(matrix, threshold = 0):
    """
        Return the network (as a networkx data structure) defined by matrix.
        It is possible to specify a threshold that will disregard all the
        edges below this threshold when creating the network
    """
    G = nx.Graph()
    N = matrix.shape[0]
    G.add_nodes_from(list(range(N)))
    G.add_weighted_edges_from([(i,j,1.0*matrix[i][j]) for i in range(0,N) for j in range(0,i) \
                                                                   if matrix[i][j] >= threshold])
    return G

def filter_weights(matrix, ratio):
    """
        Only keep a fraction of the weights of the matrix, fraction specified
        via the threshold parameter
    """
    n = matrix.shape[0]
    filtered_matrix = np.zeros_like(matrix)
    total_weight = np.sum(matrix)
    weights_id = sorted([(matrix[i,j],i,j) for i in range(n-1) for j in range(i+1,n)], reverse=True)

    filtered_weight = 0
    for (w,i,j) in weights_id:
        filtered_weight += 2*w

        if filtered_weight > ratio*total_weight:
            break

        filtered_matrix[i,j] = w
        filtered_matrix[j,i] = w

    return filtered_matrix

# -

# ## Laplacian manipulation

# +
def get_laplacian_matrix(matrix, threshold = 0):
    """
        Return the laplacian matrix of the network defined by matrix.
        It is possible to specify a threshold that will disregard all the
        edges below this threshold when creating the network
    """

    L = np.zeros_like(matrix)

    assert(L.shape[0] == L.shape[1]), "The network should be encoded by a square matrix"

    N = L.shape[0]

    # Small one liner trick to get rid of the diagonal and apply the threshold.
    A = np.multiply(np.where(matrix >= threshold, matrix , np.zeros_like(matrix)), np.ones((N,N)) - np.eye(N))

    # Use the definition of L as L = D - A
    L = np.eye(N) * np.array([np.sum(A[i,:]) for i in range(N)]) - A

    return L

def get_stabilized_laplacian_matrix(matrix, threshold = 0, ratio=1.0, alpha = 1e-2):
    """
        Prefered function to get the laplacian matrix of a network encoded by a
        matrix. Can apply a ratio filtering of the weights of the matrix (i.e.
        only keep ratio of the total weight) or / and apply a threshold. The
        alpha parameter corrects a numerical issue.
    """
    L = get_laplacian_matrix(filter_weights(matrix, ratio = ratio), threshold)

    n = L.shape[0]

    # Add a small epsilon to the diagonal of the Laplacian matrix. It should
    # make the computation more stable. Idea from https://github.com/Hermina/GOT

    return L + alpha * np.eye(n)

# -

# Small tests to assess consistency of the different functions.
assert np.allclose(nx.linalg.laplacian_matrix(get_network(connectivity_matrices[controls[0]])).toarray(), get_laplacian_matrix(connectivity_matrices[controls[0]])), "Inconsistency in the Laplacian functions"
assert np.allclose(nx.linalg.laplacian_matrix(get_network(connectivity_matrices[patients[0]], threshold=10)).toarray(), get_stabilized_laplacian_matrix(connectivity_matrices[patients[0]], alpha=0, threshold=10)), "Inconsistency in the Laplacian functions (test 2)"

# ## Wasserstein distance functions

# +
# Those are dictionaries used to speed up computations, by avoiding multiple
# computations of the same objects.
precomp_dict_net = {}
precomp_dict_cov = {}

def wasserstein_networks(A, B, threshold = 0, threshold2 = None, ratio = 1.0):
    """
            Returns the Wasserstein distance (in the sense of the GOT paper)
            between the two networks encoded by A and B. It is possible to
            apply a threshold to the network obtained via both matrices (and
            even a different one per matrices)
    """

    if threshold2 is None:
        threshold2 = threshold

    assert(A.shape == B.shape), "compared networks should have the \
                                                        same number of nodes"

    L1dag = None
    L1dag_sqrt = None
    L2dag = None
    L2dag_sqrt = None

    if (A.tobytes(), threshold, ratio) not in precomp_dict_net:
        L1 = get_stabilized_laplacian_matrix(A, threshold = threshold, ratio = ratio)
        L1dag = np.real(scp.linalg.pinv(L1))
        L1dag_sqrt = np.real(scp.linalg.sqrtm(L1dag))

        precomp_dict_net[(A.tobytes(), threshold,ratio)] = (L1dag, L1dag_sqrt)
    else:
        L1dag, L1dag_sqrt = precomp_dict_net[(A.tobytes(), threshold, ratio)]

    if (B.tobytes(), threshold2, ratio) not in precomp_dict_net:
        L2 = get_stabilized_laplacian_matrix(B, threshold = threshold2, ratio = ratio)
        L2dag = np.real(scp.linalg.pinv(L2))
        L2dag_sqrt = np.real(scp.linalg.sqrtm(L2dag))

        precomp_dict_net[(B.tobytes(), threshold2, ratio)] = (L2dag, L2dag_sqrt)

    else :
        L2dag, L2dag_sqrt = precomp_dict_net[(B.tobytes(), threshold2, ratio)]

    W = np.trace(L1dag) + np.trace(L2dag) - 2*np.trace(np.real(scp.linalg.sqrtm(L1dag_sqrt @ L2dag @ L1dag_sqrt)))

    return W

def wasserstein_network_covariance(A, cov, threshold = 0, ratio = 1.0):
    """
            Return the Wasserstein distance (in the sense of the GOT paper)
            between the the network encoded by a matrix A, and the normal
            distribution whose covariance matrix is given by cov.
            It is possible to apply a threshold to the network obtained via
            matrix A.
    """

    assert(A.shape == cov.shape), "compared networks should have the \
                                                        same number of nodes"

    N = A.shape[0]

    L1dag = None
    L1dag_sqrt = None
    L2dag = None
    L2dag_sqrt = None

    if (A.tobytes(), threshold, ratio) not in precomp_dict_net:
        L1 = get_stabilized_laplacian_matrix(A, threshold = threshold, ratio = ratio)
        L1dag = np.real(scp.linalg.pinv(L1))
        L1dag_sqrt = np.real(scp.linalg.sqrtm(L1dag))

        precomp_dict_net[(A.tobytes(), threshold ,ratio)] = (L1dag, L1dag_sqrt)
    else:
        L1dag, L1dag_sqrt = precomp_dict_net[(A.tobytes(), threshold ,ratio)]

    if cov.tobytes() not in precomp_dict_cov:
        L2dag = cov
        L2dag_sqrt = np.real(scp.linalg.sqrtm(L2dag))

        precomp_dict_cov[cov.tobytes()] = L2dag_sqrt

    else:
        L2dag = cov
        L2dag_sqrt = precomp_dict_cov[cov.tobytes()]

    W = np.trace(L1dag) + np.trace(L2dag) - 2*np.trace(np.real(scp.linalg.sqrtm(L1dag_sqrt @ L2dag @ L1dag_sqrt)))
    return W

def clip_remove(array, threshold):
    """
        Small helper function to clip an array by removing values above a
        certain threshold
    """
    r = array[array < threshold]
    print("Kept {} out of {} ({} %)".format(r.shape[0], array.shape[0], r.shape[0]/array.shape[0] * 100))
    return r

def barycenter_iteration(cov_list, iteration_count = 10):
    """
        Return the covariance matrix of the barycenter of the distributions
        whose covariances are given as cov_list ; iteration_count controls the
        number of iteration made by the iterative algorithm.
    """
    covariance_count = len(cov_list)

    size = cov_list[0].shape[0]
    curr_cov = np.eye(size)

    for iteration in range(iteration_count):
        curr_cov_sqrt = np.real(scp.linalg.sqrtm(curr_cov))
        curr_cov_pinv_sqrt = scp.linalg.pinv(curr_cov_sqrt)
        big_sum = np.zeros_like(curr_cov)
        for sigma in cov_list:
            big_sum += 1/covariance_count * np.real(scp.linalg.sqrtm(curr_cov_sqrt @ sigma @ curr_cov_sqrt))
        big_sum = big_sum @ big_sum
        curr_cov = curr_cov_pinv_sqrt @ big_sum @ curr_cov_pinv_sqrt

    return curr_cov

# Some functions to print a barycenter covariance as an approximate connectivity matrix
def cov_to_laplacian(cov, alpha = 1e-2):
    n = cov.shape[0]
    return scp.linalg.pinv(cov) - alpha * np.eye(n)

def laplacian_to_connectivity_matrix(laplacian):
    n = laplacian.shape[0]
    return np.multiply(laplacian, -(np.ones_like(laplacian) - 2*np.eye(n)))
# -

# ## Compute a distance matrix (distance subjects to controls)

# ## Compute a barycenter with subsample of controls

def subsample_controls(controls, sample_size, ratio = 1.0):
    """
        Return a random draw of sample_size controls, whose name are taken from
        the controls array, as well as the covariance matrices of their
        associated distribution. Ratio is used to control the sparsity of the
        networks (see get_stabilized_laplacian_matrix for its exact effect)
    """
    controls_count = len(controls)

    random_controls_sample = random.sample(range(0, controls_count),sample_size)
    controls_barycenter = np.array(controls)[random_controls_sample]

    pseudo_inverses = [scp.linalg.pinv(get_stabilized_laplacian_matrix(connectivity_matrices[subject],ratio=ratio)) for subject in controls_barycenter]

    return (controls_barycenter, pseudo_inverses)

def do_some_folds(fold_count = 50, control_sample_size = 10, ratios = [1.0]):
    """
        Perform fold_count random folds. Each fold consists in: subsampling the
        control group to randomly select control_sample_size controls, whose
        barycenter is computed. Then all remaining controls + each patients are
        compared to the barycenter. The results of multiple folds are aggregated
        into distance_dict and distance_fold_count (the average distance for
        subject s to the barycenters is thus

                        distance_dict[s] / distance_fold_count[s]

        ratios can be used to perform tests on multiple sparsity ratios at a
        time
    """
    for ratio in ratios:
        true_folds = 0
        print("Folding with sparsity ratio {:.2f}".format(ratio))

        pickle_file_path = pickle_folder + "/dist_and_fold_count_r{:.2f}_k{}.pickle".format(ratio, control_sample_size)
        print("Will save or update {}".format(pickle_file_path))

        distance_dict = {s : 0.0 for s in patients + controls}
        distance_fold_count = {s:0 for s in patients + controls}

        for fold in range(fold_count):
            if fold % 10 == 0: print("Fold {}".format(fold))

            controls_sample, pinvs = subsample_controls(controls, control_sample_size, ratio = ratio)

            barycenter_cov = barycenter_iteration(pinvs)

            # list of all the other subjects (non sampled controls + patients)
            subjects = patients + [c for c in controls if c not in controls_sample]
            current_fold_distance_dict = {subject:wasserstein_network_covariance(connectivity_matrices[subject], barycenter_cov, ratio = ratio) for subject in subjects}

            for subject, d in current_fold_distance_dict.items():
                distance_fold_count[subject] += 1
                distance_dict[subject] += d

            true_folds += 1

        # create or update the pickled data
        if os.path.exists(pickle_file_path) and not recompute_pickle:
            with open(pickle_file_path, 'rb') as filehandler:
                loaded_dist_dict, loaded_dist_count, loaded_fold_count = pickle.load(filehandler)

                for s in patients + controls:
                    distance_dict[s] += loaded_dist_dict[s]
                    distance_fold_count[s] += loaded_dist_count[s]

                print("{} ({} folds) will be updated with {} more folds".format(pickle_file_path, loaded_fold_count, fold_count))
                true_folds += loaded_fold_count

        with open(pickle_file_path, 'wb') as filehandler:
            print("Saving {}".format(pickle_file_path))
            pickle.dump((distance_dict, distance_fold_count, true_folds), filehandler)

        # plot the results

        # first, obtain the true distance lists from the distance_dict and
        # distance_fold_count dictionnaries
        dist_list = []
        for s in patients:
            dist_list.append(distance_dict[s] / distance_fold_count[s])

        for c in controls:
            dist_list.append(distance_dict[c] / distance_fold_count[c])

        #print(scp.stats.ttest_ind(dist_list[:patients_count], dist_list[patients_count:], equal_var=True))

        #threshold = 0.0025
        #print("Plot thresholded at {}".format(threshold))
        plt.hist(dist_list[:patients_count], stacked=True, density=True, alpha=0.3, label='Patients')
        plt.hist(dist_list[patients_count:], stacked=True, density=True, alpha = 0.3, label='Controls')
        sns.kdeplot(dist_list[:patients_count], shade=False, color='C0', alpha=0.3);
        sns.kdeplot(dist_list[patients_count:], shade=False, color='C1', alpha=0.3);
        plt.xlabel("Distance to barycenter ({:.2f})".format(ratio))
        plt.legend()
        #import tikzplotlib
        #tikzplotlib.save("figures_generation/dist_to_barycenter_full.tex".format(ratio), wrap=False)

        plt.show()

do_some_folds()

import numpy,import scipy,from scipy,import networkx


# ---
# jupyter:
#   jupytext:
#     formats: py:percent,ipynb
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.4.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # 1 Initial Approach with Dynamic Programming
# This notebook contains some intial trials of hands-on, basic reinforcement learning algorithms to familiarize with the topic and eventually arrive at the bare frame of the environment for ramp-up scenarios. The set-up follows the discourse in [Artificial Intelligence: Reinforcement Learning in Python](https://www.udemy.com/share/1013kmBEQbdF1XTXQ=/) by Lazy Programmer Inc roughly, altough the environment conforms with OpenAI's toolkit for RL, [Gym](https://gym.openai.com/).

# %%
import gym
import numpy as np
import pandas as pd
from gym import spaces
from IPython.core.display import HTML, display

# %% [markdown]
# ## Review helper functions

# %% [markdown]
# - With `display_side_by_side()`, tables are displayed side by side to ease interpretations and save space as the table index is similar.
# - Future value at each possible combination of state and episode is generated and displayed with `pivot_v()`
# - The policy at each possible combination of state and episode is generated and displayed with `pivot_p()`

# %%
def display_side_by_side(dfs: list, captions: list):
    output = ""
    combined = dict(zip(captions, dfs))
    for caption, df in combined.items():
        output += (
            df.style.set_table_attributes("style='display:inline'")
            .set_caption(caption)
            ._repr_html_()
        )
        output += "\xa0\xa0\xa0"
    display(HTML(output))

def pivot_v(V):
    v_table = pd.Series(V).reset_index()
    v_table.columns = ["Episode", "Action", "Value"]
    v_table.sort_values(by=["Episode", "Action"], inplace=True)
    v_table_p = v_table.pivot(index="Episode", columns="Action")
    return v_table, v_table_p

def pivot_p(policy):
    p_table = pd.Series(policy).reset_index()
    p_table.columns = ["Episode", "State", "Policy_S"]
    p_table.sort_values(by=["Episode"], inplace=True)
    p_table = p_table.pivot(index="Episode", columns="State")
    p_table = p_table.assign(demand=demand)
    return p_table

# %% [markdown]
# ## Definitions of constants
#
# All constants needed for the environment include
# - the **action list**, a discrete set of actions to move between states
# - the **legal changes**, a constraining factor which actions may follow upon which state
# - **demand**, for now just a list which is the trajectory that needs to be followed during an episode, only the choice of actions is open to the agent

# %%
action_list = [
    "OPERATE",  # 0
    "PREPARE",  # 1
    "PARK",  # 2
    "STORE",  # 3
]

legal_changes = {
    0: {0, 1, 2},  # from OPERATE to OPERATE, PREPARE, or PARK
    1: {1, 0, 2},  # from PREPARE to PREPARE, OPERATE, or PARK
    2: {2, 1, 3},  # from PARK to PARK, PREPARE, or STORE
    3: {3, 2},  # from STORE to STORE or PARK
}
demand = [
    0,
    0,
    0,
    0,
    0,
    0,
    0,  # week 1
    0,
    0,
    0,
    2,
    0,
    3,
    1,  # week 2
    0,
    0,
    60,
    25,
    32,
    87,
    56,  # week 3
    92,
    83,
    29,
    40,
    86,
    70,
    45,  # week 4
]

# %% [markdown]
# ## Environment Definition
# The environment is in accordance with Gym. It's a sequence of X days, on each day a single action can be performed. That's the entire logic encoded in the environment!

# %%
class RampupEnv(gym.Env):
    def __init__(self, horizon=10, verbose=False):
        super(RampupEnv, self).__init__()

        self.fleet_size = 1
        self.horizon = horizon

        self.state_time = 0  # time component of state
        self.state_status = np.zeros(self.horizon)  # status component of state

        # action_space: The Space object corresponding to valid actions
        # observation_space: The Space object corresponding to valid observations
        self.action_space = spaces.Discrete(len(action_list))
        self.observation_space = spaces.Tuple(
            (spaces.Discrete(self.fleet_size), spaces.Discrete(self.horizon))
        )

        self.total_reward = 0
        self.verbose = verbose

    def translate_action(self, action):
        if action == 0:
            action_description = "OPERATE"
            if demand[self.state_time] > 80:
                reward = 15000
            else:
                reward = -3000
        elif action == 1:
            action_description = "PREPARE"
            reward = -2000
        elif action == 2:
            action_description = "PARK"
            reward = -1000
        elif action == 3:
            action_description = "STORE"
            reward = -500

        if self.verbose:
            print(f"Action {action} leads to state {action_description}")

        return action_description, reward

    def step(self, action):
        # Increment time component of state
        self.state_time += 1

        action_description, reward = self.translate_action(action)

        # Adjust the status component of state
        self.state_status[self.state_time] = action

        obs = self.state_status[0 : self.state_time]
        done = bool(self.state_time == self.horizon - 1)
        info = {}

        if self.verbose:
            print(f"Action taken: {action_description}")

        self.total_reward += reward

        return obs, reward, done, info

    def undo_step(self, o, r, d, i):
        self.state_time -= 1
        self.total_reward -= r

        return o, r, d, i

    def render(self):
        if self.verbose:
            print(f"Reward so far: {self.total_reward}")

    def reset(self):
        # Initialize the agent in the first field
        self.state_time = 0
        self.total_reward = 0

    def set_state(self, state):
        state_time, state_status = state
        self.state_time = state_time
        self.state_status[self.state_time] = state_status

# %% [markdown]
# ## Random Walk Episode
# The length of the episode is determined by `horizon`, actions are taken at random. We observe, that games mostly lead to a negative total reward, thus it takes a policy better than a random policy to perform well.

# %%
horizon = 7 * 4
env = RampupEnv(horizon=horizon)

for step in range(horizon):
    action = env.action_space.sample()
    obs, reward, done, info = env.step(action)
    env.render()
    if done:
        print(
            f"Random walk episode is complete.\nTrajectory:\n{obs}\nTotal reward: {env.total_reward}"
        )
        break

# %% [markdown]
# ## Value Iteration
# All possible state transitions (transitions between possible combinations of state time and state actions) are iteratively evaluated for their value, until the improvement in value converges to a defined threshold. It takes 23 iterations for the sample environment with the provided data. The resulting table shows the value for each state-action.

# %%
h = 7 * 4  # time horizon
env = RampupEnv(horizon=h, verbose=False)
all_states = set(
    [(s, a) for s in range(len(env.state_status)) for a in range(env.action_space.n)]
)

SMALL_ENOUGH = 1e-3  # threshold for convergence
GAMMA = 1.0  # discount factor

V = {}
for state in all_states:
    V[state] = 0

biggest_change = 0
iters = 0

while True:
    iters += 1
    biggest_change = 0
    for state in all_states:
        old_v = V[state]
        env.set_state(state)
        state_time, state_status = state
        if env.state_time < env.horizon - 1:
            new_v = 0  # answer is accumulated
            p_a = 1.0 / len(legal_changes[state_status])  # equal probability
            for a in legal_changes[state_status]:
                env.set_state(state)
                o, r, d, i = env.step(a)
                new_v += p_a * (r + GAMMA * V[(env.state_time, a)])
                o, r, d, i = env.undo_step(o, r, d, i)
            V[state] = new_v
            biggest_change = max(biggest_change, np.abs(old_v - V[state]))

        if env.verbose:
            print(f"State: {state}, Old V: {old_v}")

    if biggest_change < SMALL_ENOUGH:
        print(f"Done in {iters} iterations")
        v_table, v_table_p = pivot_v(V)
        display(v_table_p)
        break

# %% [markdown]
# ## Policy Iteration
# Value iteration stops as the improvement of value converges, policy iteration then adds a policy improvement step. In the case of the example, there is no improvement visible, but we do get the policy for the values from value iterations inside of the attached policy improvement step.

# %%
h = 7 * 4  # time horizon
env = RampupEnv(horizon=h, verbose=False)
all_states = set(
    [(s, a) for s in range(len(env.state_status)) for a in range(env.action_space.n)]
)

SMALL_ENOUGH = 1e-3
GAMMA = 1

policy = {}
for state in all_states:
    policy[state] = env.action_space.sample()

V = {}
for state in all_states:
    V[state] = 0

iters1, iters2 = 0, 0
rep1, rep2 = "", ""

while True:
    iters1 = 0
    # Policy Evaluation Step
    while True:
        iters1 += 1
        biggest_change = 0
        for state in all_states:
            old_v = V[state]
            env.set_state(state)
            state_time, state_status = state
            if env.state_time < env.horizon - 1:
                new_v = 0  # answer is accumulated
                p_a = 1.0 / len(legal_changes[state_status])  # equal probability
                for a in legal_changes[state_status]:
                    env.set_state(state)
                    o, r, d, i = env.step(a)
                    new_v += p_a * (r + GAMMA * V[(env.state_time, a)])
                    o, r, d, i = env.undo_step(o, r, d, i)
                V[state] = new_v
                biggest_change = max(biggest_change, np.abs(old_v - V[state]))

            if env.verbose:
                print(f"State: {state}, Old V: {old_v}")

        if biggest_change < SMALL_ENOUGH:
            rep1 += f"{str(iters1)}, "
            iters2 += 1
            break

    # Policy Improvement Step
    is_policy_converged = True
    for state in all_states:
        for state in policy:
            old_a = policy[state]
            new_a = None
            best_value = float("-inf")
            env.set_state(state)
            state_time, state_status = state
            if env.state_time < env.horizon - 1:
                # loop through all possible actions to find the best current action
                for a in legal_changes[state_status]:
                    env.set_state(state)
                    o, r, d, i = env.step(a)
                    v = r + GAMMA * V[(env.state_time, a)]
                    # print(v)
                    if v > best_value:
                        best_value = v
                        new_a = a
                policy[state] = new_a
                if new_a != old_a:
                    is_policy_converged = False
    if is_policy_converged:
        print(
            f"Policy evaluation and improvement steps were entered {iters2} times. \
Policy evaluation iterated {rep1[:-2]} times."
        )
        v_table, v_table_p = pivot_v(V)
        p_table = pivot_p(policy)
        display_side_by_side([v_table_p, p_table], ["Value Table", "Policy Table"])
        break

# %%


r"""
Facade Sets
"""
#*****************************************************************************
#  Copyright (C) 2010 Nicolas M. Thiery <nthiery at users.sf.net>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#                  http://www.gnu.org/licenses/
#******************************************************************************

from sage.misc.cachefunc import cached_method
from sage.categories.category_singleton import Category_singleton
from sage.categories.category import Category

class FacadeSets(Category_singleton):
    r"""
    The category of facade sets

    A *facade set* is a parent ``P`` whose elements actually belong to
    some other parent::

        sage: P = Sets().example(); P
        Set of prime numbers (basic implementation)
        sage: p = Sets().example().an_element(); p
        47
        sage: p in P
        True
        sage: p.parent()
        Integer Ring

    Typical use cases include modeling a subset of an existing
    parent::

        sage: Sets().Facades().example()
        An example of facade set: the monoid of positive integers

    or the union of several parents::

        sage: Sets().Facades().example("union")
        An example of a facade set: the integers completed by +-infinity

    or endowing a parent with more (or less!) structure::

        sage: Posets().example("facade")
        An example of a facade poset: the positive integers ordered by divisibility

    Let us consider one of the examples above in detail: the partially ordered
    set `P` of positive integers w.r.t. divisibility order. There are two
    options for representing its elements:

     1. as plain integers
     2. as integers, modified to be aware that their parent is `P`

    The advantage of 1. is that one needs not to do conversions back and
    forth between `P` and `\ZZ`. The disadvantage is that this
    introduces an ambiguity when writing `2 < 3`::

        sage: 2 < 3
        True

    To raise this ambiguity, one needs to explicitely specify the order
    as in `2 <_P 3`::

        sage: P = Posets().example("facade")
        sage: P.lt(2,3)
        False

    In short `P` being a facade parent is one of the programmatic
    counterpart (with e.g. coercions) of the usual mathematical idiom:
    "for ease of notation, we identify an element of `P` with the
    corresponding integer". Too many identifications lead to
    confusion; the lack thereof leads to heavy, if not obfuscated,
    notations. Finding the right balance is an art, and even though
    there are common guidelines, it is ultimately up to the writer to
    choose which identifications to do. This is no different in code.

    .. seealso::

       ::

        sage: Sets().example("facade")
        Set of prime numbers (facade implementation)
        sage: Sets().example("inherits")
        Set of prime numbers
        sage: Sets().example("wrapper")
        Set of prime numbers (wrapper implementation)

    .. rubric:: Specifications

    A parent which is a facade must either:

    - call :meth:`Parent.__init__` using the ``facade`` parameter to
      specify a parent, or tuple thereof.
    - overload the method :meth:`~FacadeSets.ParentMethods.facade_for`.

    .. note:: the concept of facade parents was originally introduced
       in the computer algebra system MuPAD.

    TESTS:

    Check that multiple categories initialisation works (:trac:`13801`)::

        sage: class A(Parent):
        ...     def __init__(self):
        ...         Parent.__init__(self, category=(FiniteEnumeratedSets(),Monoids()), facade=True)
        sage: a = A()
    """

    @cached_method
    def super_categories(self):
        r"""
        Returns the super categories of ``self``, as per
        :meth:`Category.super_categories`.

        EXAMPLES::

            sage: Sets().Facades().super_categories()
            [Category of sets]
        """
        from sage.categories.sets_cat import Sets
        return [Sets()]

    def example(self, choice='subset'):
        r"""
        Returns an example of facade set, as per
        :meth:`Category.example()
        <sage.categories.category.Category.example>`.

        INPUT:

        - ``choice`` -- 'union' or 'subset' (default: 'subset').

        EXAMPLES::

            sage: Sets().Facades().example()
            An example of facade set: the monoid of positive integers
            sage: Sets().Facades().example(choice='union')
            An example of a facade set: the integers completed by +-infinity
            sage: Sets().Facades().example(choice='subset')
            An example of facade set: the monoid of positive integers
        """
        import sage.categories.examples.facade_sets as examples
        if choice == "union":
            return examples.IntegersCompletion()
        elif choice == 'subset':
            return examples.PositiveIntegerMonoid()
        else:
            raise TypeError, "choice should be 'union' or 'subset'"

    class ParentMethods:

        def _element_constructor_(self, element):
            """
            Coerce ``element`` into ``self``

            INPUT:

            - ``element`` -- any object

            This default implementation returns ``element`` if
            ``self`` is a facade for ``parent(element)`. Otherwise it
            attempts in turn to coerce ``element`` into each parent
            ``self`` is a facade for.

            This implementation is only valid for a facade parent
            which models the full union of the parents it is a facade
            for. Other facade parents should redefine
            :meth:`element_constructor` appropriately.

            EXAMPLES::

                sage: S = Sets().Facades().example("union"); S
                An example of a facade set: the integers completed by +-infinity
                sage: S(1)
                1
                sage: S(1/2)
                Traceback (most recent call last):
                ...
                ValueError: Can't coerce `1/2` in any parent `An example of a facade set: the integers completed by +-infinity` is a facade for
                sage: S(2/1)
                2
                sage: S(2/1).parent()
                Integer Ring
                sage: S(int(1))
                1
                sage: S(int(1)).parent()
                Integer Ring

            Facade parents that model strict subsets should redefine
            :meth:`element_constructor`::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S(-1)
                Traceback (most recent call last):
                ...
                ValueError: %s should be positive
            """
            if self.is_parent_of(element):
                return element
            else:
                parents = self.facade_for()
                if parents is True:
                    return NotImplementedError
                for parent in self.facade_for():
                    try:
                        return parent(element)
                    except Exception:
                        pass
            raise ValueError, "Can't coerce `%s` in any parent `%s` is a facade for"%(element, self)

        def facade_for(self):
            """
            Returns the parents this set is a facade for

            This default implementation assumes that ``self`` has
            an attribute ``_facade_for``, typically initialized by
            :meth:`Parent.__init__`. If the attribute is not present, the method
            raises a NotImplementedError.

            EXAMPLES::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S.facade_for()
                (Integer Ring,)

            Check that :trac:`13801` is corrected::

                sage: class A(Parent):
                ...     def __init__(self):
                ...         Parent.__init__(self, category=Sets(), facade=True)
                sage: a = A()
                sage: a.facade_for()
                Traceback (most recent call last):
                ...
                NotImplementedError: this parent did not specify which parents it is a facade for
            """
            try:
                return self._facade_for
            except AttributeError:
                raise NotImplementedError("this parent did not specify which parents it is a facade for")

        def is_parent_of(self, element):
            """
            Returns whether ``self`` is the parent of ``element``

            INPUT:

            - ``element`` -- any object

            Since ``self`` is a facade domain, this actually tests
            whether the parent of ``element`` is any of the parent
            ``self`` is a facade for.

            EXAMPLES::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S.is_parent_of(1)
                True
                sage: S.is_parent_of(1/2)
                False

            This method differs from :meth:`__contains__` in two
            ways.  First, this does not take into account the fact
            that ``self`` may be a strict subset of the parent(s)
            it is a facade for::

                sage: -1 in S, S.is_parent_of(-1)
                (False, True)

            Furthermore, there is no coercion attempted::

                sage: int(1) in S, S.is_parent_of(int(1))
                (True, False)

            .. warning::

               this implementation does not handle facade parents of facade
               parents. Is this a feature we want generically?
            """
            parents = self.facade_for()
            if parents is True:
                return True
            from sage.structure.element import parent
            return parent(element) in parents

        def __contains__(self, element):
            """
            Membership testing

            Returns whether ``element`` is in one of the parents
            ``self`` is a facade for.

            .. warning:: this default implementation is currently
            overriden by :meth:`Parent.__contains__`.

            EXAMPLES::

                sage: S = Sets().Facades().example("union"); S
                An example of a facade set: the integers completed by +-infinity
                sage: 1 in S, -5 in S, oo in S, -oo in S, int(1) in S, 2/1 in S
                (True, True, True, True, True, True)
                sage: 1/2 in S, "bla" in S
                (False, False)
            """
            return any(element in parent for parent in self.facade_for())

        def _an_element_(self):
            """
            Try to return an element of ``self``, as per
            :meth:`Sets.ParentMethods.an_element`.

            For each parent ``self`` is a facade for, this default
            implementation tries the method ``an_element`` until it finds an
            element in ``self``. If none is found raise a
            ``NotImplementedError``.

            EXAMPLES::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S.an_element()
                1
            """
            for parent in self.facade_for():
                x = parent.an_element()
                if x in self:
                    return x
            raise NotImplementedError