
import numpy as np
import pandas as pd
import pandas.io.data as web
import matplotlib.pyplot as plt
import timeit
# import seaborn as sns
from scipy import stats
from pandas.io.pytables import HDFStore
from datetime import datetime
from Data import Data
from Portfolio import Portfolio
import rs

def bestwindow(p):
    """
    Experiment for finding optimal window size
    1. Select portfolio p with initial window size
    2. Observe returns for range of investment horizons
    3. Compare sample distributions
    """
    windowReturns = p.getReturns()

    # find cdf for window returns    
    # f = lambda ih: stats.ks_2samp(windowReturns, p.getPortfolioFutureReturns(ih))[1] # second is p-value
    params = stats.johnsonsu.fit(windowReturns)
    f = lambda ih: stats.kstest(p.getPortfolioFutureReturns(ih), 'johnsonsu', args=params )[1] # second is p-value

    ihss = np.linspace(10, p.window, 200)
    ihs = pd.Series(ihss, index=ihss )
    pValues = ihs.apply(f)

    # print ihs, pValues
    
    # df = pd.DataFrame({'ihs':ihs, 'pValues':pValues})
    # df.plot(x='ihs', y='pValues')
    # # sns.lmplot('ihs', 'pValues', df)
    # # plt.plot(ihs, pValues)
    # plt.show()
    return pValues



def bestWindowSimulation(p):
    df = pd.DataFrame()
    for d in pd.date_range(p.date, periods=12*6, freq='M'):
        p.setup(d)
        pValues = bestwindow(p)
        df[d] = pValues        
    df.index = pValues.index
    return df

def particlesCount(p):
    """ Calculate utility function over iterations for different particle count """
    p.setup() # inisial state with equal weights    
    uList = pd.DataFrame() #data frame for list of utility functions
    for popSize in [10, 20, 30, 50, 70, 100, 130]: #range(10, 101, 20):
        print popSize
        meanList = pd.DataFrame()
        for i in range(3):
            meanList[i] = p.optimize('pso', opts={'pop':popSize, 'iterations':500})

        uList[popSize] = meanList.mean(axis=1)

    return uList

def pso_c(p):
    i = np.linspace(1.0, 1.8, 5)
    
    counter = 0
    ac1 = []
    ac2 = []
    ay = []
    x = 0
    for c1 in i:
        for c2 in i:
            print c1, c2
            meanList = []
            for c in range(3):
                u = p.optimize('pso', opts={"pop": 20, 'c1':c1, 'c2':c2, 'w':0.7, 'iterations': 1000})
                meanList.append(u.count())            
                x+=1
                print x, "-----------------------------------------------------"
                                    
            ac1.append(c1)
            ac2.append(c2)
            ay.append(pd.Series(meanList).mean())    

    df = pd.DataFrame({'c1':ac1, 'c2':ac2, 'y':ay}, columns=['c1', 'c2', 'y'])    
    #plot        
    return df

def saturation(p):
    it = []
    fit = []
    for i in range(10):
        u = p.optimize('pso', opts={"pop": 20, 'c1':2.0, 'c2':2.0, 'w':0.7})
        fit.append(p.sortino())
        it.append(u.count())
    return pd.DataFrame({'it':pd.Series(it), 'fit':pd.Series(fit)}, columns=['it', 'fit'])


def optimize_d_p(di,pi, p):
    meanList = []
    meanUList = []
    for c in range(3):
        t = datetime.now()                
        p[di].optimize("pso", opts={"pop": pi, 'iterations':1000})
        meanList.append((datetime.now() - t).microseconds)
        meanUList.append(p[di].sortino())
    return pd.Series(meanList).mean(), pd.Series(meanUList).mean()


def particles_dimensions(data):
    
    # dimensions array
    d = np.array(range(5,36,5))

    # particles array
    p = np.array(range(10,71,10)) #range(10,20,3)


    # create dict of portfolios
    portfolios = {di: Portfolio(pd.Series(data.assetList), data, date=datetime(2012, 1, 1), window=365*2, maxAssets=di) for di in d }

    d, p = np.meshgrid(d, p)

    
    optimize_d_p_vectorized = np.vectorize(optimize_d_p, excluded=['p'])
    
    t, u = optimize_d_p_vectorized(d,p, portfolios)
    t = np.nan_to_num(t)
    u = np.nan_to_num(u)
            
    return d,p,t,u

def anneal_vs_pso(p):
    uList = pd.DataFrame()
    meanList = pd.DataFrame()
    for i in range(3):
        meanList[i] = p.optimize('pso')

    uList['pso'] = meanList.mean(axis=1)

    meanList = pd.DataFrame()
    for i in range(3):
        meanList[i] = p.optimize('anneal')

    uList['anneal'] = meanList.mean(axis=1)
    plt.xlabel('iteration')
    plt.ylabel('utility')
    uList.plot()
    plt.show()

    return uList




    
