"""MC1-P2: Optimize a portfolio."""

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import datetime as dt
from util import get_data, plot_data
import math as m
import scipy.optimize as spo

# This is the function that will be tested by the autograder
# The student must update this code to properly implement the functionality
def optimize_portfolio(sd=dt.datetime(2008,1,1), ed=dt.datetime(2009,1,1), \
    syms=['GOOG','AAPL','GLD','XOM'], gen_plot=False):

    # Read in adjusted closing prices for given symbols, date range
    dates = pd.date_range(sd, ed)
    prices_all = get_data(syms, dates, sd, ed)  # automatically adds SPY
    prices = prices_all[syms]  # only portfolio symbols
    prices_SPY = prices_all['SPY']  # only SPY, for comparison later

    # find the allocations for the optimal portfolio
    sv = 1000000
    normalized_prices = prices / prices.ix[0, :]
    x0 = np.array([0.2, 0.2, 0.3, 0.3, 0.0])
    optimal_allocs = spo.minimize(f, x0, args=(normalized_prices, sv),
                        method='SLSQP', options={'disp': True},
                        bounds=tuple((0, 1) for i in range(0, x0.size)),
                        constraints = ({'type': 'eq', 'fun': \
                        lambda inputs: 1.0 - np.sum(inputs)})
                        )
    allocs = optimal_allocs.x
    port_val, cr, adr, sddr, sr = calc_portfolio_stats(allocs, normalized_prices, sv)

    # Compare daily portfolio value with SPY using a normalized plot
    if gen_plot:
        df_temp = pd.concat([port_val, prices_SPY], keys=['Portfolio', 'SPY'], axis=1)
        df_temp = df_temp / df_temp.ix[0, :]
        df_temp.plot()
        plt.show()
        pass

    return allocs, cr, adr, sddr, sr

def f(x, normalized_prices, sv):

    port_val, cr, adr, sddr, sr = calc_portfolio_stats(x, normalized_prices, sv)

    # Return Sharpe Ratio
    # Multiply by -1 because we are minimizing this function
    return sr * -1

def calc_portfolio_stats(allocs, normalized_prices, sv):
    # Calculate position values and sum for portfolio values
    position_values = allocs * normalized_prices * sv
    port_val = position_values.sum(axis=1)

    # Get portfolio statistics
    portfolio_returns = (port_val / port_val.shift(1) - 1).fillna(0)
    cr = (portfolio_returns + 1).cumprod()[-1] - 1
    adr = portfolio_returns.mean()
    sddr = portfolio_returns.std()
    sr = ((adr - 0) / sddr) * m.sqrt(252)

    return port_val, cr, adr, sddr, sr

def test_code():
    # This function WILL NOT be called by the auto grader
    # Do not assume that any variables defined here are available to your function/code
    # It is only here to help you set up and test your code

    # Define input parameters
    # Note that ALL of these values will be set to different values by
    # the autograder!

    start_date = dt.datetime(2009,1,1)
    end_date = dt.datetime(2010,1,1)
    symbols = ['GOOG', 'AAPL', 'GLD', 'XOM', 'IBM']

    # Assess the portfolio
    allocations, cr, adr, sddr, sr = optimize_portfolio(sd = start_date, ed = end_date,\
        syms = symbols, \
        gen_plot = False)

    # Print statistics
    print "Start Date:", start_date
    print "End Date:", end_date
    print "Symbols:", symbols
    print "Allocations:", allocations
    print "Sharpe Ratio:", sr
    print "Volatility (stdev of daily returns):", sddr
    print "Average Daily Return:", adr
    print "Cumulative Return:", cr

if __name__ == "__main__":
    # This code WILL NOT be called by the auto grader
    # Do not assume that it will be called
    test_code()
