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

import os

import numpy as np
import matplotlib.pyplot as plt

import pywt

ecg = np.load(os.path.join('data', 'ecg.npy'))

data1 = np.concatenate((np.arange(1, 400),
                        np.arange(398, 600),
                        np.arange(601, 1024)))
x = np.linspace(0.082, 2.128, num=1024)[::-1]
data2 = np.sin(40 * np.log(x)) * np.sign((np.log(x)))

mode = pywt.MODES.sp1

def plot_signal_decomp(data, w, title):
    """Decompose and plot a signal S.

    S = An + Dn + Dn-1 + ... + D1
    """
    w = pywt.Wavelet(w)
    a = data
    ca = []
    cd = []
    for i in range(5):
        (a, d) = pywt.dwt(a, w, mode)
        ca.append(a)
        cd.append(d)

    rec_a = []
    rec_d = []

    for i, coeff in enumerate(ca):
        coeff_list = [coeff, None] + [None] * i
        rec_a.append(pywt.waverec(coeff_list, w))

    for i, coeff in enumerate(cd):
        coeff_list = [None, coeff] + [None] * i
        rec_d.append(pywt.waverec(coeff_list, w))

    fig = plt.figure()
    ax_main = fig.add_subplot(len(rec_a) + 1, 1, 1)
    ax_main.set_title(title)
    ax_main.plot(data)
    ax_main.set_xlim(0, len(data) - 1)

    for i, y in enumerate(rec_a):
        ax = fig.add_subplot(len(rec_a) + 1, 2, 3 + i * 2)
        ax.plot(y, 'r')
        ax.set_xlim(0, len(y) - 1)
        ax.set_ylabel("A%d" % (i + 1))

    for i, y in enumerate(rec_d):
        ax = fig.add_subplot(len(rec_d) + 1, 2, 4 + i * 2)
        ax.plot(y, 'g')
        ax.set_xlim(0, len(y) - 1)
        ax.set_ylabel("D%d" % (i + 1))

plot_signal_decomp(data1, 'coif5', "DWT: Signal irregularity")
plot_signal_decomp(data2, 'sym5', "DWT: Frequency and phase change - Symmlets5")
plot_signal_decomp(ecg, 'sym5', "DWT: Ecg sample - Symmlets5")

plt.show()


import pytest

import numpy
import pandas
import pandas.util.testing as tm

def pytest_addoption(parser):
    parser.addoption("--skip-slow", action="store_true",
                     help="skip slow tests")
    parser.addoption("--skip-network", action="store_true",
                     help="skip network tests")
    parser.addoption("--run-high-memory", action="store_true",
                     help="run high memory tests")
    parser.addoption("--only-slow", action="store_true",
                     help="run only slow tests")

def pytest_runtest_setup(item):
    if 'slow' in item.keywords and item.config.getoption("--skip-slow"):
        pytest.skip("skipping due to --skip-slow")

    if 'slow' not in item.keywords and item.config.getoption("--only-slow"):
        pytest.skip("skipping due to --only-slow")

    if 'network' in item.keywords and item.config.getoption("--skip-network"):
        pytest.skip("skipping due to --skip-network")

    if 'high_memory' in item.keywords and not item.config.getoption(
            "--run-high-memory"):
        pytest.skip(
            "skipping high memory test since --run-high-memory was not set")

# Configurations for all tests and all test modules

@pytest.fixture(autouse=True)
def configure_tests():
    pandas.set_option('chained_assignment', 'raise')

# For running doctests: make np and pd names available

@pytest.fixture(autouse=True)
def add_imports(doctest_namespace):
    doctest_namespace['np'] = numpy
    doctest_namespace['pd'] = pandas

@pytest.fixture(params=['bsr', 'coo', 'csc', 'csr', 'dia', 'dok', 'lil'])
def spmatrix(request):
    tm._skip_if_no_scipy()
    from scipy import sparse
    return getattr(sparse, request.param + '_matrix')

@pytest.fixture
def ip():
    """
    Get an instance of IPython.InteractiveShell.

    Will raise a skip if IPython is not installed.
    """

    pytest.importorskip('IPython', minversion="6.0.0")
    from IPython.core.interactiveshell import InteractiveShell
    return InteractiveShell()


import matplotlib
import paths
matplotlib.rc_file(paths.pcpath('pubfiguresrc'))
import os
import pylab as pl
from astropy import table
from paths import analysispath
import numpy as np
from astropy import coordinates
from astropy import units as u
import heating

pcfittable = table.Table.read(os.path.join(analysispath,
                                         'fitted_line_parameters_Chi2Constraints.ipac'),
                            format='ascii.ipac')

lolim = pcfittable['tmax1sig_chi2'] > 340
maps = np.char.startswith(pcfittable['Source_Name'], 'Map')
ok = ~np.isnan(pcfittable['tmin1sig_chi2']) & (pcfittable['width'] < 40) & (pcfittable['h2coratio321303']/pcfittable['eh2coratio321303'] > 5) & pcfittable['is_good'].astype('bool')
flags = {'is_map': maps,
         'is_lolim': lolim,
         'is_ok': ok}
# Don't plot these for now...
pcfittable = pcfittable[(~lolim) & ok]
maps = np.char.startswith(pcfittable['Source_Name'], 'Map')
lolim_conservative = pcfittable['tmax1sig_chi2'] > 150

fig4 = pl.figure(4)
fig4.clf()
ax = fig4.add_subplot(1,3,1)
ax.errorbar(pcfittable['temperature_chi2'], pcfittable['density_chi2'],
            yerr=[pcfittable['density_chi2']-pcfittable['dmin1sig_chi2'],
                  pcfittable['dmax1sig_chi2']-pcfittable['density_chi2']],
            xerr=[pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'],
                  pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2']],
            linestyle='none', marker='s', linewidth=1, alpha=0.5)
ax2 = fig4.add_subplot(1,3,2)
# Don't do this any more: it relies on having the RADEX fits, which we don't.
#ax2.errorbar(pcfittable['temperature_chi2'], pcfittable['temperature'],
#             yerr=pcfittable['etemperature'],
#             xerr=[pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'],
#                   pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2']],
#             linestyle='none', marker='s', linewidth=1, alpha=0.5)
ax2.plot([0,300],[0,300],'k--',linewidth=2,alpha=0.5)

fig5 = pl.figure(5)
fig5.clf()
ax5 = fig5.gca()
ax5.errorbar(coordinates.Angle(pcfittable['GLON']*u.deg).wrap_at(180*u.deg).value[maps],
             pcfittable['temperature_chi2'][maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[maps]],
             capsize=0, markeredgecolor='none',
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
ax5.set_ylim(0,150)
ax5.set_ylabel("Temperature (K)")
ax5.set_xlabel("Galactic Longitude ($^{\\circ}$)")
fig5.savefig(paths.fpath('chi2_temperature_vs_glon_byfield.pdf'),
                         bbox_inches='tight')
ax5.errorbar(coordinates.Angle(pcfittable['GLON']*u.deg).wrap_at(180*u.deg).value[~maps],
             pcfittable['temperature_chi2'][~maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[~maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[~maps]],
             capsize=0, markeredgecolor='none',
             linestyle='none', marker='s', linewidth=1, alpha=0.5)
fig5.savefig(paths.fpath('chi2_temperature_vs_glon_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig6 = pl.figure(6)
fig6.clf()
ax6 = fig6.gca()
mask = maps&~lolim_conservative
ax6.errorbar(pcfittable['higaldusttem'][mask],
             pcfittable['temperature_chi2'][mask],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[mask],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[mask]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r', capsize=0)
ax6.plot([15,30],[15,30],'k--')
mask = maps&lolim_conservative
ax6.plot(pcfittable['higaldusttem'][mask],
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=10,
         markeredgecolor='none',
         color='r',
         alpha=0.5,
         linestyle='none')
ax6.set_xlabel("HiGal Dust Temperature (K)")
ax6.set_ylabel("H$_2$CO Temperature (K)")
ax6.set_ylim(0,200)
ax6.set_xlim(15,30)
fig6.savefig(paths.fpath('chi2_temperature_vs_higaltemperature_byfield.pdf'),
                         bbox_inches='tight')
mask = (~maps)&(~lolim_conservative)
ax6.errorbar(pcfittable['higaldusttem'][mask],
             pcfittable['temperature_chi2'][mask],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[mask],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[mask]],
             capsize=0,
             markeredgecolor='none',
             markersize=10,
             linestyle='none', marker='s', linewidth=0.5, alpha=0.5, color='b')

mask = (~maps)&lolim_conservative
ax6.plot(pcfittable['higaldusttem'][mask],
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=10,
         markeredgecolor='none',
         color='b',
         alpha=0.5,
         linestyle='none')

ax6.set_ylim(10,150)
ax6.set_xlim(15,30)
fig6.savefig(paths.fpath('chi2_temperature_vs_higaltemperature_fieldsandsources_notitle.pdf'),
                         bbox_inches='tight')
ax6.set_title("Hand-selected regions")
fig6.savefig(paths.fpath('chi2_temperature_vs_higaltemperature_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig7 = pl.figure(7)
fig7.clf()
ax7 = fig7.gca()
mask = maps&~lolim_conservative
ax7.errorbar(pcfittable['width'][mask]*(8*np.log(2))**0.5,
             pcfittable['temperature_chi2'][mask],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[mask],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[mask]],
             capsize=0,
             markersize=10,
             markeredgecolor='none',
             linestyle='none', marker='s', linewidth=0.5, alpha=0.6, color='r')
mask = maps&lolim_conservative
ax7.plot(pcfittable['width'][mask]*(8*np.log(2))**0.5,
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=10,
         markeredgecolor='none',
         color='r',
         alpha=0.4,
         linestyle='none')

linewidths = np.linspace(0,pcfittable['width'].max())*u.km/u.s
ax7.plot(linewidths*2.35, [heating.tkin_all(10**4*u.cm**-3, sigma, 10*u.pc,
                                            5*u.km/u.s/u.pc, 30*u.K)
                           for sigma in linewidths],
        linestyle='--', color='k', label='$n=10^4$ cm$^{-3}$', zorder=-5)
ax7.plot(linewidths*2.35, [heating.tkin_all(10**4*u.cm**-3, sigma, 10*u.pc,
                                            1*u.km/u.s/u.pc, 30*u.K)
                           for sigma in linewidths],
        linestyle='--', color='r', label='$n=10^4$ cm$^{-3}$, $dv/dr=1$', zorder=-5, linewidth=2, alpha=0.5)
ax7.plot(linewidths*2.35, [heating.tkin_all(10**4*u.cm**-3, sigma, 20*u.pc,
                                            5*u.km/u.s/u.pc, 30*u.K)
                           for sigma in linewidths],
         linestyle='--', color='b', label='$n=10^4$ cm$^{-3}$, $L=20$ pc', zorder=-5, alpha=0.5, linewidth=2)

ax7.plot(linewidths*2.35, [heating.tkin_all(10**5*u.cm**-3, sigma, 10*u.pc,
                                            5*u.km/u.s/u.pc, 30*u.K)
                           for sigma in linewidths],
         linestyle=':', color='k', label='$n=10^5$ cm$^{-3}$', zorder=-5)
ax7.plot(linewidths*2.35, [heating.tkin_all(10**6*u.cm**-3, sigma, 10*u.pc,
                                            5*u.km/u.s/u.pc, 30*u.K)
                           for sigma in linewidths],
        linestyle='-.', color='k', label='$n=10^6$ cm$^{-3}$', zorder=-5)
ax7.plot(linewidths*2.35, [heating.tkin_all(10**5*u.cm**-3, sigma, 10*u.pc,
                                            5*u.km/u.s/u.pc, 30*u.K, crir=1e-15*u.s**-1)
                           for sigma in linewidths],
         linestyle='-', color='g', label='$n=10^5$ cm$^{-3}$, $\zeta_{CR}=10^{-15}$ s$^{-1}$', zorder=-10, alpha=0.25, linewidth=4)
ax7.plot(linewidths*2.35, [heating.tkin_all(10**5*u.cm**-3, sigma, 10*u.pc,
                                            5*u.km/u.s/u.pc, 30*u.K, crir=1e-14*u.s**-1)
                           for sigma in linewidths],
         linestyle=':', color='purple', label='$n=10^5$ cm$^{-3}$, $\zeta_{CR}=10^{-14}$ s$^{-1}$', zorder=-10, alpha=0.25, linewidth=4)

box = ax7.get_position()
ax7.set_position([box.x0, box.y0, box.width * 0.7, box.height])
ax7.legend(loc='center left', fontsize=16, bbox_to_anchor=(1.0, 0.75))
ax7.set_xlabel("Line FWHM (km s$^{-1}$)")
ax7.set_ylabel("Temperature (K)")
ax7.set_ylim(10,150)
fig7.savefig(paths.fpath('chi2_temperature_vs_linewidth_byfield.pdf'),
                         bbox_inches='tight')

mask = (~maps)&(~lolim_conservative)
ax7.errorbar(pcfittable['width'][mask]*(8*np.log(2))**0.5,
             pcfittable['temperature_chi2'][mask],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[mask],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[mask]],
             capsize=0,
             markeredgecolor='none',
             markersize=10,
             linestyle='none', marker='s', linewidth=0.5, alpha=0.6, color='b')

mask = (~maps)&lolim_conservative
ax7.plot(pcfittable['width'][mask]*(8*np.log(2))**0.5,
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=10,
         markeredgecolor='none',
         color='b',
         alpha=0.4,
         linestyle='none')

ax7.set_ylim(10,150)
fig7.savefig(paths.fpath('chi2_temperature_vs_linewidth_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig8 = pl.figure(8)
fig8.clf()
ax8 = fig8.gca()
ax8.errorbar(pcfittable['ampH2CO'][maps],
             pcfittable['temperature_chi2'][maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[maps]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
ax8.set_xlabel("H2CO Peak Amplitude")
ax8.set_ylabel("Temperature (K)")
fig8.savefig(paths.fpath('chi2_temperature_vs_h2coamp_byfield.pdf'),
                         bbox_inches='tight')
ax8.errorbar(pcfittable['ampH2CO'][~maps],
             pcfittable['temperature_chi2'][~maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[~maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[~maps]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='b')
fig8.savefig(paths.fpath('chi2_temperature_vs_h2coamp_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig9 = pl.figure(9)
fig9.clf()
ax9 = fig9.gca()
ax9.set_xscale('log')
ax9.errorbar(pcfittable['higalcolumndens'][maps],
             pcfittable['temperature_chi2'][maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[maps]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
ax9.set_xlabel("Hi-Gal Fitted Column Density")
ax9.set_ylabel("Temperature (K)")
fig9.savefig(paths.fpath('chi2_temperature_vs_higalcolumn_byfield.pdf'),
                         bbox_inches='tight')
ax9.errorbar(pcfittable['higalcolumndens'][~maps],
             pcfittable['temperature_chi2'][~maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[~maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[~maps]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='b')
fig9.savefig(paths.fpath('chi2_temperature_vs_higalcolumn_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig10 = pl.figure(10)
fig10.clf()
ax10 = fig10.gca()
ax10.errorbar(pcfittable['width'][maps]*(8*np.log(2))**0.5,
             pcfittable['h2coratio321303'][maps],
             yerr=pcfittable['eh2coratio321303'][maps],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
ax10.set_xlabel("Line FWHM (km s$^{-1}$)")
ax10.set_ylabel("Ratio 321/303")
fig10.savefig(paths.fpath('ratio_vs_linewidth_byfield.pdf'),
                         bbox_inches='tight')
ax10.errorbar(pcfittable['width'][~maps]*(8*np.log(2))**0.5,
              pcfittable['h2coratio321303'][~maps],
              yerr=pcfittable['eh2coratio321303'][~maps],
              linestyle='none', marker='s', linewidth=1, alpha=0.5, color='b')
fig10.savefig(paths.fpath('ratio_vs_linewidth_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig11 = pl.figure(11)
fig11.clf()
ax11 = fig11.gca()
ax11.errorbar(pcfittable['higaldusttem'][maps],
             pcfittable['h2coratio321303'][maps],
             yerr=pcfittable['eh2coratio321303'][maps],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
ax11.set_ylim(0,200)
ax11.set_xlim(15,30)
ax11.set_xlabel("HiGal Fitted Temperature")
ax11.set_ylabel("Ratio 321/303")
fig11.savefig(paths.fpath('ratio_vs_higaltemperature_byfield.pdf'),
                         bbox_inches='tight')
ax11.errorbar(pcfittable['higaldusttem'][~maps],
              pcfittable['h2coratio321303'][~maps],
              yerr=pcfittable['eh2coratio321303'][~maps],
              linestyle='none', marker='s', linewidth=1, alpha=0.5, color='b')
fig11.savefig(paths.fpath('ratio_vs_higaltemperature_fieldsandsources.pdf'),
                         bbox_inches='tight')

# RADEX fitting has been removed
#fig12 = pl.figure(12)
#fig12.clf()
#ax = fig12.add_subplot(1,1,1)
#ax.errorbar(pcfittable['temperature_chi2'], pcfittable['temperature'],
#            yerr=pcfittable['etemperature'],
#            xerr=[pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'],
#                  pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2']],
#            linestyle='none', marker='s', linewidth=1, alpha=0.5)
#ax.plot([0,300],[0,300],'k--',linewidth=2,alpha=0.5)
#ax.set_title("DEBUG: RADEX+pyspeckit-fitted temperature vs. $\\chi^2$ temperature")
#ax.set_xlabel("$\\chi^2$ Temperature")
#ax.set_ylabel("RADEX+pyspeckit Temperature")
#ax.axis([0,350,0,350])

fig13 = pl.figure(13)
fig13.clf()
ax13 = fig13.gca()
ax13.errorbar(pcfittable['area'][maps],
              pcfittable['temperature_chi2'][maps],
              yerr=[pcfittable['temperature_chi2'][maps]-pcfittable['tmin1sig_chi2'][maps],
                    pcfittable['tmax1sig_chi2'][maps]-pcfittable['temperature_chi2'][maps]],
              linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
ax13.set_xlabel("Area (square degrees)")
ax13.set_ylabel("Temperature (K)")
ax13.set_xscale('log')
fig13.savefig(paths.fpath('temperature_vs_area_byfield.pdf'),
                         bbox_inches='tight')
ax13.errorbar(pcfittable['area'][~maps],
              pcfittable['temperature_chi2'][~maps],
              yerr=[pcfittable['temperature_chi2'][~maps]-pcfittable['tmin1sig_chi2'][~maps],
                    pcfittable['tmax1sig_chi2'][~maps]-pcfittable['temperature_chi2'][~maps]],
              linestyle='none', marker='s', linewidth=1, alpha=0.5, color='b')
fig13.savefig(paths.fpath('temperature_vs_area_fieldsandsources.pdf'),
                         bbox_inches='tight')

fig14 = pl.figure(14)
fig14.clf()
ax14 = fig14.gca()
ax14.errorbar(pcfittable['higalcolumndens'][maps],
             pcfittable['temperature_chi2'][maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[maps]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='r')
#ax14.plot([15,30],[15,30],'k--')
ax14.set_xlabel("HiGal Fitted Column Density")
ax14.set_ylabel("Temperature (K)")
fig14.savefig(paths.fpath('chi2_temperature_vs_higaldustcol_byfield.pdf'),
                         bbox_inches='tight')
ax14.errorbar(pcfittable['higalcolumndens'][~maps],
             pcfittable['temperature_chi2'][~maps],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[~maps],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[~maps]],
             linestyle='none', marker='s', linewidth=1, alpha=0.5, color='b')
fig14.savefig(paths.fpath('chi2_temperature_vs_higaldustcol_fieldsandsources.pdf'),
                         bbox_inches='tight')

# pcfittable[np.abs(pcfittable['temperature_chi2']-pcfittable['higaldusttem'])/pcfittable['higaldusttem'] < 1.5].pprint()

fig15 = pl.figure(15)
fig15.clf()
ax15 = fig15.gca()
mask = maps&~lolim_conservative
ax15.errorbar(pcfittable['tkin_turb'][mask],
              pcfittable['temperature_chi2'][mask],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[mask],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[mask]],
             capsize=0,
             markersize=10,
             markeredgecolor='none',
             linestyle='none', marker='s', linewidth=0.5, alpha=0.6, color='r')
mask = maps&lolim_conservative
ax15.plot(pcfittable['tkin_turb'][mask],
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=10,
         markeredgecolor='none',
         color='r',
         alpha=0.4,
         linestyle='none')

mask = (maps) & (~lolim_conservative) & ((pcfittable['tmin1sig_chi2'] > pcfittable['tkin_turb']) | (pcfittable['tmax1sig_chi2'] < pcfittable['tkin_turb']))
ax15.plot(pcfittable['tkin_turb'][mask],
         pcfittable['temperature_chi2'][mask],
         marker='s',
         markersize=15,
         markeredgecolor='r',
         markerfacecolor='none',
         markeredgewidth=0.5,
         alpha=0.4,
         linestyle='none')
mask = (maps) & (lolim_conservative) & ((pcfittable['tmin1sig_chi2'] > pcfittable['tkin_turb']))
ax15.plot(pcfittable['tkin_turb'][mask],
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=15,
         markeredgecolor='r',
         markerfacecolor='none',
         markeredgewidth=0.5,
         alpha=0.4,
         linestyle='none')

# Sources with T_predicted >> T_measured
#high_badpredictions = (pcfittable['tkin_turb'] > pcfittable['tmax1sig_chi2'])&(~lolim_conservative)
#high_badpredictions = (pcfittable['tkin_turb'] > 120)&(~lolim_conservative)
#for row,is_map in zip(pcfittable[high_badpredictions], maps[high_badpredictions]):
#    xy = np.array((row['tkin_turb'], row['temperature_chi2']))
#    ax15.annotate("{0}_{1}".format(row['Source_Name'], row['ComponentID']),
#                  xy,
#                  xytext=xy-(15, 7),
#                  color='r' if is_map else 'b'
#                 )

ax15.plot([0,200], [0,200], 'k--', alpha=0.5, zorder=-5)
ax15.set_xlabel("Turbulence-driven Temperature (K)")
ax15.set_ylabel("H$_2$CO Temperature (K)")
ax15.set_ylim(10,150)
ax15.set_xlim(10,180)
fig15.savefig(paths.fpath('chi2_temperature_vs_turbulenttemperature_byfield.pdf'),
                         bbox_inches='tight')

mask = (~maps)&(~lolim_conservative)
ax15.errorbar(pcfittable['tkin_turb'][mask],
             pcfittable['temperature_chi2'][mask],
             yerr=[(pcfittable['temperature_chi2']-pcfittable['tmin1sig_chi2'])[mask],
                   (pcfittable['tmax1sig_chi2']-pcfittable['temperature_chi2'])[mask]],
             capsize=0,
             markeredgecolor='none',
             markersize=10,
             linestyle='none', marker='s', linewidth=0.5, alpha=0.6, color='b')

mask = (~maps)&lolim_conservative
ax15.plot(pcfittable['tkin_turb'][mask],
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=10,
         markeredgecolor='none',
         color='b',
         alpha=0.4,
         linestyle='none')

mask = (~maps) & (~lolim_conservative) & ((pcfittable['tmin1sig_chi2'] > pcfittable['tkin_turb']) | (pcfittable['tmax1sig_chi2'] < pcfittable['tkin_turb']))
ax15.plot(pcfittable['tkin_turb'][mask],
         pcfittable['temperature_chi2'][mask],
         marker='s',
         markersize=15,
         markeredgecolor='b',
         markerfacecolor='none',
         markeredgewidth=0.5,
         alpha=0.4,
         linestyle='none')
mask = (~maps) & (lolim_conservative) & ((pcfittable['tmin1sig_chi2'] > pcfittable['tkin_turb']))
ax15.plot(pcfittable['tkin_turb'][mask],
         pcfittable['tmin1sig_chi2'][mask],
         marker='^',
         markersize=15,
         markeredgecolor='b',
         markerfacecolor='none',
         markeredgewidth=0.5,
         alpha=0.4,
         linestyle='none')

ax15.set_ylim(10,150)
fig15.savefig(paths.fpath('chi2_temperature_vs_turbulenttemperature_fieldsandsources_notitle.pdf'),
                         bbox_inches='tight')
ax15.set_title("Hand-selected regions")
fig15.savefig(paths.fpath('chi2_temperature_vs_turbulenttemperature_fieldsandsources.pdf'),
                         bbox_inches='tight')


import numpy
import opengm
import vigra
import matplotlib.pyplot as plt
import matplotlib.cm as cm

gradScale       = 0.1
energyNotEqual  = 0.2
sigma=0.2
resizeFactor=2

img=vigra.impex.readImage('lena.bmp')
shape=img.shape
imgLab=vigra.colors.transform_RGB2Lab(img)
shape=(shape[0]*resizeFactor,shape[1]*resizeFactor)
imgLab=vigra.sampling.resize(imgLab, shape,order=3)

gradMag=vigra.filters.gaussianGradientMagnitude(imgLab,gradScale)

unaries=numpy.zeros([shape[0],shape[1],2])
unaries[:,:,1]=numpy.exp(-1.0*gradMag[:,:,0]*sigma)
unaries[:,:,0]=1.0-unaries[:,:,1]
regularizer=opengm.PottsFunction([2,2],0.0,energyNotEqual)

gm=opengm.grid2d2Order(unaries=unaries,regularizer=regularizer,order='numpy',operator='adder')
inf=opengm.inference.GraphCut(gm)
inf.infer()
argmin=inf.arg().reshape(shape[0:2])

plt.figure(1)

ax=plt.subplot(2,1,1)
plt.imshow(unaries[:,:,1].T, interpolation="nearest")
plt.set_cmap(cm.copper)
plt.colorbar()
ax.set_title('costs / unaries label=1')

ax=plt.subplot(2,1,2)
plt.imshow(argmin.T, interpolation="nearest")
plt.colorbar()
ax.set_title('argmin')

plt.show()


from __future__ import division
"""

 backend_wxagg.py

 A wxPython backend for Agg.  This uses the GUI widgets written by
 Jeremy O'Donoghue (jeremy@o-donoghue.com) and the Agg backend by John
 Hunter (jdhunter@ace.bsd.uchicago.edu)

 Copyright (C) 2003-5 Jeremy O'Donoghue, John Hunter, Illinois Institute of
 Technology

 License: This work is licensed under the matplotlib license( PSF
 compatible). A copy should be included with this source code.

"""

import wx
import matplotlib
from matplotlib.figure import Figure

from backend_agg import FigureCanvasAgg
import backend_wx
from backend_wx import FigureManager, FigureManagerWx, FigureCanvasWx, \
    FigureFrameWx, DEBUG_MSG, NavigationToolbar2Wx, error_msg_wx, \
    draw_if_interactive, show, Toolbar, backend_version

class FigureFrameWxAgg(FigureFrameWx):
    def get_canvas(self, fig):
        return FigureCanvasWxAgg(self, -1, fig)

    def _get_toolbar(self, statbar):
        if matplotlib.rcParams['toolbar']=='classic':
            toolbar = NavigationToolbarWx(self.canvas, True)
        elif matplotlib.rcParams['toolbar']=='toolbar2':
            toolbar = NavigationToolbar2WxAgg(self.canvas)
            toolbar.set_status_bar(statbar)
        else:
            toolbar = None
        return toolbar

class FigureCanvasWxAgg(FigureCanvasAgg, FigureCanvasWx):
    """
    The FigureCanvas contains the figure and does event handling.

    In the wxPython backend, it is derived from wxPanel, and (usually)
    lives inside a frame instantiated by a FigureManagerWx. The parent
    window probably implements a wxSizer to control the displayed
    control size - but we give a hint as to our preferred minimum
    size.
    """

    def draw(self, drawDC=None):
        """
        Render the figure using agg.
        """
        DEBUG_MSG("draw()", 1, self)
        FigureCanvasAgg.draw(self)

        self.bitmap = _convert_agg_to_wx_bitmap(self.get_renderer(), None)
        self._isDrawn = True
        self.gui_repaint(drawDC=drawDC)

    def blit(self, bbox=None):
        """
        Transfer the region of the agg buffer defined by bbox to the display.
        If bbox is None, the entire buffer is transferred.
        """
        if bbox is None:
            self.bitmap = _convert_agg_to_wx_bitmap(self.get_renderer(), None)
            self.gui_repaint()
            return

        l, b, w, h = bbox.bounds
        r = l + w
        t = b + h
        x = int(l)
        y = int(self.bitmap.GetHeight() - t)

        srcBmp = _convert_agg_to_wx_bitmap(self.get_renderer(), None)
        srcDC = wx.MemoryDC()
        srcDC.SelectObject(srcBmp)

        destDC = wx.MemoryDC()
        destDC.SelectObject(self.bitmap)

        destDC.BeginDrawing()
        destDC.Blit(x, y, int(w), int(h), srcDC, x, y)
        destDC.EndDrawing()

        destDC.SelectObject(wx.NullBitmap)
        srcDC.SelectObject(wx.NullBitmap)
        self.gui_repaint()

    filetypes = FigureCanvasAgg.filetypes

    def print_figure(self, filename, *args, **kwargs):
        # Use pure Agg renderer to draw
        FigureCanvasAgg.print_figure(self, filename, *args, **kwargs)
        # Restore the current view; this is needed because the
        # artist contains methods rely on particular attributes
        # of the rendered figure for determining things like
        # bounding boxes.
        if self._isDrawn:
            self.draw()

class NavigationToolbar2WxAgg(NavigationToolbar2Wx):
    def get_canvas(self, frame, fig):
        return FigureCanvasWxAgg(frame, -1, fig)

def new_figure_manager(num, *args, **kwargs):
    """
    Create a new figure manager instance
    """
    # in order to expose the Figure constructor to the pylab
    # interface we need to create the figure here
    DEBUG_MSG("new_figure_manager()", 3, None)
    backend_wx._create_wx_app()

    FigureClass = kwargs.pop('FigureClass', Figure)
    fig = FigureClass(*args, **kwargs)
    frame = FigureFrameWxAgg(num, fig)
    figmgr = frame.get_figure_manager()
    if matplotlib.is_interactive():
        figmgr.frame.Show()
    return figmgr

#
# agg/wxPython image conversion functions (wxPython <= 2.6)
#

def _py_convert_agg_to_wx_image(agg, bbox):
    """
    Convert the region of the agg buffer bounded by bbox to a wx.Image.  If
    bbox is None, the entire buffer is converted.

    Note: agg must be a backend_agg.RendererAgg instance.
    """
    image = wx.EmptyImage(int(agg.width), int(agg.height))
    image.SetData(agg.tostring_rgb())

    if bbox is None:
        # agg => rgb -> image
        return image
    else:
        # agg => rgb -> image => bitmap => clipped bitmap => image
        return wx.ImageFromBitmap(_clipped_image_as_bitmap(image, bbox))

def _py_convert_agg_to_wx_bitmap(agg, bbox):
    """
    Convert the region of the agg buffer bounded by bbox to a wx.Bitmap.  If
    bbox is None, the entire buffer is converted.

    Note: agg must be a backend_agg.RendererAgg instance.
    """
    if bbox is None:
        # agg => rgb -> image => bitmap
        return wx.BitmapFromImage(_py_convert_agg_to_wx_image(agg, None))
    else:
        # agg => rgb -> image => bitmap => clipped bitmap
        return _clipped_image_as_bitmap(
            _py_convert_agg_to_wx_image(agg, None),
            bbox)

def _clipped_image_as_bitmap(image, bbox):
    """
    Convert the region of a wx.Image bounded by bbox to a wx.Bitmap.
    """
    l, b, width, height = bbox.get_bounds()
    r = l + width
    t = b + height

    srcBmp = wx.BitmapFromImage(image)
    srcDC = wx.MemoryDC()
    srcDC.SelectObject(srcBmp)

    destBmp = wx.EmptyBitmap(int(width), int(height))
    destDC = wx.MemoryDC()
    destDC.SelectObject(destBmp)

    destDC.BeginDrawing()
    x = int(l)
    y = int(image.GetHeight() - t)
    destDC.Blit(0, 0, int(width), int(height), srcDC, x, y)
    destDC.EndDrawing()

    srcDC.SelectObject(wx.NullBitmap)
    destDC.SelectObject(wx.NullBitmap)

    return destBmp

#
# agg/wxPython image conversion functions (wxPython >= 2.8)
#

def _py_WX28_convert_agg_to_wx_image(agg, bbox):
    """
    Convert the region of the agg buffer bounded by bbox to a wx.Image.  If
    bbox is None, the entire buffer is converted.

    Note: agg must be a backend_agg.RendererAgg instance.
    """
    if bbox is None:
        # agg => rgb -> image
        image = wx.EmptyImage(int(agg.width), int(agg.height))
        image.SetData(agg.tostring_rgb())
        return image
    else:
        # agg => rgba buffer -> bitmap => clipped bitmap => image
        return wx.ImageFromBitmap(_WX28_clipped_agg_as_bitmap(agg, bbox))

def _py_WX28_convert_agg_to_wx_bitmap(agg, bbox):
    """
    Convert the region of the agg buffer bounded by bbox to a wx.Bitmap.  If
    bbox is None, the entire buffer is converted.

    Note: agg must be a backend_agg.RendererAgg instance.
    """
    if bbox is None:
        # agg => rgba buffer -> bitmap
        return wx.BitmapFromBufferRGBA(int(agg.width), int(agg.height),
            agg.buffer_rgba(0, 0))
    else:
        # agg => rgba buffer -> bitmap => clipped bitmap
        return _WX28_clipped_agg_as_bitmap(agg, bbox)

def _WX28_clipped_agg_as_bitmap(agg, bbox):
    """
    Convert the region of a the agg buffer bounded by bbox to a wx.Bitmap.

    Note: agg must be a backend_agg.RendererAgg instance.
    """
    l, b, width, height = bbox.get_bounds()
    r = l + width
    t = b + height

    srcBmp = wx.BitmapFromBufferRGBA(int(agg.width), int(agg.height),
        agg.buffer_rgba(0, 0))
    srcDC = wx.MemoryDC()
    srcDC.SelectObject(srcBmp)

    destBmp = wx.EmptyBitmap(int(width), int(height))
    destDC = wx.MemoryDC()
    destDC.SelectObject(destBmp)

    destDC.BeginDrawing()
    x = int(l)
    y = int(int(agg.height) - t)
    destDC.Blit(0, 0, int(width), int(height), srcDC, x, y)
    destDC.EndDrawing()

    srcDC.SelectObject(wx.NullBitmap)
    destDC.SelectObject(wx.NullBitmap)

    return destBmp

def _use_accelerator(state):
    """
    Enable or disable the WXAgg accelerator, if it is present and is also
    compatible with whatever version of wxPython is in use.
    """
    global _convert_agg_to_wx_image
    global _convert_agg_to_wx_bitmap

    if getattr(wx, '__version__', '0.0')[0:3] < '2.8':
        # wxPython < 2.8, so use the C++ accelerator or the Python routines
        if state and _wxagg is not None:
            _convert_agg_to_wx_image  = _wxagg.convert_agg_to_wx_image
            _convert_agg_to_wx_bitmap = _wxagg.convert_agg_to_wx_bitmap
        else:
            _convert_agg_to_wx_image  = _py_convert_agg_to_wx_image
            _convert_agg_to_wx_bitmap = _py_convert_agg_to_wx_bitmap
    else:
        # wxPython >= 2.8, so use the accelerated Python routines
        _convert_agg_to_wx_image  = _py_WX28_convert_agg_to_wx_image
        _convert_agg_to_wx_bitmap = _py_WX28_convert_agg_to_wx_bitmap

# try to load the WXAgg accelerator
try:
    import _wxagg
except ImportError:
    _wxagg = None

# if it's present, use it
_use_accelerator(True)


# 可以自己import我们平台支持的第三方python模块，比如pandas、numpy等。
import talib
import numpy as np
import math
import pandas

# 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
def init(context):
    context.s1 = "000001.XSHE"

    # 使用MACD需要设置长短均线和macd平均线的参数
    context.SHORTPERIOD = 12
    context.LONGPERIOD = 26
    context.SMOOTHPERIOD = 9
    context.OBSERVATION = 100

# 你选择的证券的数据更新将会触发此段逻辑，例如日或分钟历史数据切片或者是实时数据切片更新
def handle_bar(context, bar_dict):
    # 开始编写你的主要的算法逻辑

    # bar_dict[order_book_id] 可以拿到某个证券的bar信息
    # context.portfolio 可以拿到现在的投资组合状态信息

    # 使用order_shares(id_or_ins, amount)方法进行落单

    # TODO: 开始编写你的算法吧！

    #读取历史数据，使用sma方式计算均线准确度和数据长度无关，但是在使用ema方式计算均线时建议将历史数据窗口适当放大，结果会更加准确
    prices = history(context.OBSERVATION,'1d','close')[context.s1].values

    # 用Talib计算MACD取值，得到三个时间序列数组，分别为macd,signal 和 hist
    macd,signal,hist = talib.MACD(prices,context.SHORTPERIOD,context.LONGPERIOD,context.SMOOTHPERIOD)

    plot("macd",macd[-1])
    plot("macd signal",signal[-1])

    # macd 是长短均线的差值，signal是macd的均线，使用macd策略有几种不同的方法，我们这里采用macd线突破signal线的判断方法

    # 如果macd从上往下跌破macd_signal

    if macd[-1]-signal[-1]<0 and macd[-2]-signal[-2]>0:
        # 计算现在portfolio中股票的仓位
        curPosition = context.portfolio.positions[context.s1].quantity
        #进行清仓
        if curPosition>0:
            order_target_value(context.s1,0)

    # 如果短均线从下往上突破长均线，为入场信号
    if macd[-1]-signal[-1]>0 and macd[-2]-signal[-2]<0:
    #满仓入股
        order_target_percent(context.s1, 1)


"""
====================================================================
Normal and Shrinkage Linear Discriminant Analysis for classification
====================================================================

Shows how shrinkage improves classification.
"""

from __future__ import division

import numpy as np
import matplotlib.pyplot as plt

from sklearn.datasets import make_blobs
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis

n_train = 20  # samples for training
n_test = 200  # samples for testing
n_averages = 50  # how often to repeat classification
n_features_max = 75  # maximum number of features
step = 4  # step size for the calculation

def generate_data(n_samples, n_features):
    """Generate random blob-ish data with noisy features.

    This returns an array of input data with shape `(n_samples, n_features)`
    and an array of `n_samples` target labels.

    Only one feature contains discriminative information, the other features
    contain only noise.
    """
    X, y = make_blobs(n_samples=n_samples, n_features=1, centers=[[-2], [2]])

    # add non-discriminative features
    if n_features > 1:
        X = np.hstack([X, np.random.randn(n_samples, n_features - 1)])
    return X, y

acc_clf1, acc_clf2 = [], []
n_features_range = range(1, n_features_max + 1, step)
for n_features in n_features_range:
    score_clf1, score_clf2 = 0, 0
    for _ in range(n_averages):
        X, y = generate_data(n_train, n_features)

        clf1 = LinearDiscriminantAnalysis(solver='lsqr', shrinkage='auto').fit(X, y)
        clf2 = LinearDiscriminantAnalysis(solver='lsqr', shrinkage=None).fit(X, y)

        X, y = generate_data(n_test, n_features)
        score_clf1 += clf1.score(X, y)
        score_clf2 += clf2.score(X, y)

    acc_clf1.append(score_clf1 / n_averages)
    acc_clf2.append(score_clf2 / n_averages)

features_samples_ratio = np.array(n_features_range) / n_train

plt.plot(features_samples_ratio, acc_clf1, linewidth=2,
         label="Linear Discriminant Analysis with shrinkage", color='r')
plt.plot(features_samples_ratio, acc_clf2, linewidth=2,
         label="Linear Discriminant Analysis", color='g')

plt.xlabel('n_features / n_samples')
plt.ylabel('Classification accuracy')

plt.legend(loc=1, prop={'size': 12})
plt.suptitle('Linear Discriminant Analysis vs. \
shrinkage Linear Discriminant Analysis (1 discriminative feature)')
plt.show()


"""
==========================================
One-class SVM with non-linear kernel (RBF)
==========================================

An example using a one-class SVM for novelty detection.

:ref:`One-class SVM <svm_outlier_detection>` is an unsupervised
algorithm that learns a decision function for novelty detection:
classifying new data as similar or different to the training set.
"""
print(__doc__)

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager
from sklearn import svm

xx, yy = np.meshgrid(np.linspace(-5, 5, 500), np.linspace(-5, 5, 500))
# Generate train data
X = 0.3 * np.random.randn(100, 2)
X_train = np.r_[X + 2, X - 2]
# Generate some regular novel observations
X = 0.3 * np.random.randn(20, 2)
X_test = np.r_[X + 2, X - 2]
# Generate some abnormal novel observations
X_outliers = np.random.uniform(low=-4, high=4, size=(20, 2))

# fit the model
clf = svm.OneClassSVM(nu=0.1, kernel="rbf", gamma=0.1)
clf.fit(X_train)
y_pred_train = clf.predict(X_train)
y_pred_test = clf.predict(X_test)
y_pred_outliers = clf.predict(X_outliers)
n_error_train = y_pred_train[y_pred_train == -1].size
n_error_test = y_pred_test[y_pred_test == -1].size
n_error_outliers = y_pred_outliers[y_pred_outliers == 1].size

# plot the line, the points, and the nearest vectors to the plane
Z = clf.decision_function(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)

plt.title("Novelty Detection")
plt.contourf(xx, yy, Z, levels=np.linspace(Z.min(), 0, 7), cmap=plt.cm.PuBu)
a = plt.contour(xx, yy, Z, levels=[0], linewidths=2, colors='darkred')
plt.contourf(xx, yy, Z, levels=[0, Z.max()], colors='palevioletred')

s = 40
b1 = plt.scatter(X_train[:, 0], X_train[:, 1], c='white', s=s)
b2 = plt.scatter(X_test[:, 0], X_test[:, 1], c='blueviolet', s=s)
c = plt.scatter(X_outliers[:, 0], X_outliers[:, 1], c='gold', s=s)
plt.axis('tight')
plt.xlim((-5, 5))
plt.ylim((-5, 5))
plt.legend([a.collections[0], b1, b2, c],
           ["learned frontier", "training observations",
            "new regular observations", "new abnormal observations"],
           loc="upper left",
           prop=matplotlib.font_manager.FontProperties(size=11))
plt.xlabel(
    "error train: %d/200 ; errors novel regular: %d/40 ; "
    "errors novel abnormal: %d/40"
    % (n_error_train, n_error_test, n_error_outliers))
plt.show()


# yellowbrick.text.dispersion
# Implementations of lexical dispersions for text visualization.
#
# Author:   Larry Gray
# Created:  2018-06-21 10:06
#
# Copyright (C) 2018 District Data Labs
# For license information, see LICENSE.txt
#
# ID: dispersion.py [] lwgray@gmail.com $

"""
Implementation of lexical dispersion for text visualization
"""

##########################################################################
## Imports
##########################################################################

from collections import defaultdict
import itertools

from yellowbrick.text.base import TextVisualizer
from yellowbrick.style.colors import resolve_colors
from yellowbrick.exceptions import YellowbrickValueError

import numpy as np

##########################################################################
## Dispersion Plot Visualizer
##########################################################################

class DispersionPlot(TextVisualizer):
    """
    DispersionPlotVisualizer allows for visualization of the lexical dispersion
    of words in a corpus.  Lexical dispersion is a measure of a word's
    homeogeneity across the parts of a corpus.  This plot notes the occurences
    of a word and how many words from the beginning it appears.

    Parameters
    ----------
    target_words : list
        A list of target words whose dispersion across a corpus passed at fit
	will be visualized.

    ax : matplotlib axes, default: None
        The axes to plot the figure on.

    labels : list of strings
        The names of the classes in the target, used to create a legend.
        Labels must match names of classes in sorted order.

    colors : list or tuple of colors
        Specify the colors for each individual class

    colormap : string or matplotlib cmap
        Qualitative colormap for discrete target

    ignore_case : boolean, default: False
	Specify whether input  will be case-sensitive.

    annotate_docs : boolean, default: False
        Specify whether document boundaries will be displayed.  Vertical lines
        are positioned at the end of each document.

    kwargs : dict
        Pass any additional keyword arguments to the super class.

    These parameters can be influenced later on in the visualization
    process, but can and should be set as early as possible.
    """

    # NOTE: cannot be np.nan
    NULL_CLASS = None

    def __init__(self, target_words, ax=None, colors=None, ignore_case=False,
                 annotate_docs=False, labels=None, colormap=None, **kwargs):
        super(DispersionPlot, self).__init__(ax=ax, **kwargs)

        self.labels = labels
        self.colors = colors
        self.colormap = colormap

        self.target_words = target_words
        self.ignore_case = ignore_case
        self.annotate_docs = annotate_docs

    def _compute_dispersion(self, text, y):
        self.boundaries_ = []
        offset = 0

        if y is None:
            y = itertools.repeat(None)

        for doc, target in zip(text, y):
            for word in doc:
                if self.ignore_case:
                    word = word.lower()

                # NOTE: this will find all indices if duplicate words are supplied
                # In the case that word is not in target words, any empty list is
                # returned and no data will be yielded
                offset += 1
                for y_coord in (self.indexed_words_ == word).nonzero()[0]:
                    y_coord = int(y_coord)
                    yield (offset, y_coord, target)
            if self.annotate_docs:
                self.boundaries_.append(offset)
        self.boundaries_ = np.array(self.boundaries_, dtype=int)

    def _check_missing_words(self, points):
        for index in range(len(self.indexed_words_)):
            if index in points[:,1]:
                pass
            else:
                raise YellowbrickValueError((
                    "The indexed word '{}' is not found in "
                    "this corpus"
                    ).format(self.indexed_words_[index]))

    def fit(self, X, y=None, **kwargs):
        """
        The fit method is the primary drawing input for the dispersion
        visualization.

        Parameters
        ----------
        X : list or generator
            Should be provided as a list of documents or a generator
            that yields a list of documents that contain a list of
            words in the order they appear in the document.

        y : ndarray or Series of length n
            An optional array or series of target or class values for
            instances. If this is specified, then the points will be colored
            according to their class.

        kwargs : dict
            Pass generic arguments to the drawing method

        Returns
        -------
        self : instance
            Returns the instance of the transformer/visualizer
        """

        if y is not None:
            self.classes_ = np.unique(y)
        elif y is None and self.labels is not None:
            self.classes_ = np.array([self.labels[0]])
        else:
            self.classes_ = np.array([self.NULL_CLASS])

        # Create an index (e.g. the y position) for the target words
        self.indexed_words_ = np.flip(self.target_words, axis=0)
        if self.ignore_case:
            self.indexed_words_ = np.array([w.lower() for w in self.indexed_words_])

        # Stack is used to create a 2D array from the generator
        try:
            points_target = np.stack(self._compute_dispersion(X, y))
        except ValueError:
            raise YellowbrickValueError((
                "No indexed words were found in the corpus"
                ))
        points = np.stack(zip(points_target[:,0].astype(int),
                              points_target[:,1].astype(int)))

        self.target = points_target[:,2]

        self._check_missing_words(points)

        self.draw(points, self.target)
        return self

    def draw(self, points, target=None, **kwargs):
        """
        Called from the fit method, this method creates the canvas and
        draws the plot on it.
        Parameters
        ----------
        kwargs: generic keyword arguments.
        """

        # Resolve the labels with the classes
        labels = self.labels if self.labels is not None else self.classes_
        if len(labels) != len(self.classes_):
            raise YellowbrickValueError((
                "number of supplied labels ({}) does not "
                "match the number of classes ({})"
            ).format(len(labels), len(self.classes_)))

        # Create the color mapping for the labels.
        color_values = resolve_colors(
            n_colors=len(labels), colormap=self.colormap, colors=self.color)
        colors = dict(zip(labels, color_values))

        # Transform labels into a map of class to label
        labels = dict(zip(self.classes_, labels))

        # Define boundaries with a vertical line
        if self.annotate_docs:
            for xcoords in self.boundaries_:
                self.ax.axvline(x=xcoords, color='lightgray', linestyle='dashed')

        series = defaultdict(lambda: {'x':[], 'y':[]})

        if target is not None:
            for point, t in zip(points, target):
                label = labels[t]
                series[label]['x'].append(point[0])
                series[label]['y'].append(point[1])
        else:
            label = self.classes_[0]
            for x, y in points:
                series[label]['x'].append(x)
                series[label]['y'].append(y)

        for label, points in series.items():
            self.ax.scatter(points['x'], points['y'], marker='|',
                            c=colors[label], zorder=100, label=label)

        self.ax.set_yticks(list(range(len(self.indexed_words_))))
        self.ax.set_yticklabels(self.indexed_words_)

    def finalize(self, **kwargs):
        """
        The finalize method executes any subclass-specific axes
        finalization steps. The user calls poof & poof calls finalize.
        Parameters
        ----------
        kwargs: generic keyword arguments.
        """

        self.ax.set_ylim(-1, len(self.indexed_words_))
        self.ax.set_title("Lexical Dispersion Plot")
        self.ax.set_xlabel("Word Offset")
        self.ax.grid(False)

        # Add the legend outside of the figure box.
        if not all(self.classes_ == np.array([self.NULL_CLASS])):
            box = self.ax.get_position()
            self.ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
            self.ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))

##########################################################################
## Quick Method
##########################################################################

def dispersion(words, corpus, y=None, ax=None, colors=None, colormap=None,
               labels=None, annotate_docs=False, ignore_case=False, **kwargs):
    """ Displays lexical dispersion plot for words in a corpus

    This helper function is a quick wrapper to utilize the DisperstionPlot
    Visualizer for one-off analysis

    Parameters
    ----------

    words : list
        A list of words whose dispersion will be examined within a corpus

    y : ndarray or Series of length n
        An optional array or series of target or class values for
        instances. If this is specified, then the points will be colored
        according to their class.

    corpus : list
        Should be provided as a list of documents that contain
        a list of words in the order they appear in the document.

    ax : matplotlib axes, default: None
        The axes to plot the figure on.

    labels : list of strings
        The names of the classes in the target, used to create a legend.
        Labels must match names of classes in sorted order.

    colors : list or tuple of colors
        Specify the colors for each individual class

    colormap : string or matplotlib cmap
        Qualitative colormap for discrete target

    annotate_docs : boolean, default: False
        Specify whether document boundaries will be displayed.  Vertical lines
        are positioned at the end of each document.

    ignore_case : boolean, default: False
	Specify whether input  will be case-sensitive.

    kwargs : dict
        Pass any additional keyword arguments to the super class.

    Returns
    -------
    ax: matplotlib axes
        Returns the axes that the plot was drawn on
    """

    # Instantiate the visualizer
    visualizer = DispersionPlot(
        words, ax=ax, colors=colors, colormap=colormap,
        ignore_case=ignore_case, labels=labels,
        annotate_docs=annotate_docs, **kwargs
    )

    # Fit and transform the visualizer (calls draw)
    visualizer.fit(corpus, y, **kwargs)

    # Return the axes object on the visualizer
    return visualizer.ax


import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import plot, show, xticks, xlabel, ylabel, legend, yscale, title, savefig, rcParams, figure, hist, text, bar, subplots
import Image

def file_process(filename):
    fig = plt.gcf()
    fig.set_size_inches(18.5,10.5)
    fig.savefig("/home/rob/%s.png" % filename,dpi=300)
    fig.savefig("/home/rob/%sLOWRES.png" % filename,dpi=50)

    Image.open("/home/rob/%s.png" % filename).convert('L').save("/home/rob/%s.jpg" % filename)
    Image.open("/home/rob/%sLOWRES.png" % filename).convert('L').save("/home/rob/%sLOWRES.jpg" % filename)

"""
compare:

handcrafting

bootstrapped

one shot

equal weights

market cap weights
"""

import pandas as pd
from datetime import datetime as dt

def read_ts_csv(fname, dindex="Date"):
    data=pd.read_csv(fname)
    dateindex=[dt.strptime(dx, "%d/%m/%y") for dx in list(data[dindex])]
    data.index=dateindex
    del(data[dindex])

    return data

def calc_asset_returns(rawdata, tickers):
    asset_returns=pd.concat([get_monthly_tr(tickname, rawdata) for tickname in tickers], axis=1)
    asset_returns.columns=tickers

    return asset_returns

def get_monthly_tr(tickname, rawdata):
    total_returns=rawdata[tickname+"_TR"]
    return (total_returns / total_returns.shift(1)) - 1.0

def portfolio_return(asset_returns, cash_weights):

    index_returns=asset_returns.cumsum().ffill().diff()

    cash_align = cash_weights.reindex(asset_returns.index, method="ffill")
    cash_align[np.isnan(index_returns)]=0.0
    cash_align[np.isnan(cash_align)]=0.0

    vols=pd.ewmstd(asset_returns, span=100, min_periods=1)
    riskweights=pd.DataFrame(cash_align.values / vols.values, index=vols.index)
    riskweights.columns=asset_returns.columns

    riskweights[np.isnan(riskweights)]=0.0

    def _rowfix(x):
        if all([y==0.0 for y in x]):
            return x
        sumx=sum(x)
        return [y/sumx for y in x]

    riskweights = riskweights.apply(_rowfix, axis=1)

    portfolio_returns=asset_returns*riskweights

    portfolio_returns[np.isnan(portfolio_returns)]=0.0

    portfolio_returns=portfolio_returns.sum(axis=1)

    return portfolio_returns

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

from scipy.optimize import minimize
from copy import copy
import random

def correlation_matrix(returns):
    """
    Calcs a correlation matrix using weekly returns from a pandas time series

    We use weekly returns because otherwise end of day effects, especially over time zones, give
    unrealistically low correlations

    """
    asset_index=returns.cumsum().ffill()
    asset_index=asset_index.resample('1W') ## Only want index, fill method is irrelevant
    asset_index = asset_index - asset_index.shift(1)

    return asset_index.corr().values

def create_dull_pd_matrix(dullvalue=0.0, dullname="A", startdate=pd.datetime(1970,1,1).date(), enddate=datetime.datetime.now().date(), index=None):
    """
    create a single valued pd matrix
    """
    if index is None:
        index=pd.date_range(startdate, enddate)

    dullvalue=np.array([dullvalue]*len(index))

    ans=pd.DataFrame(dullvalue, index, columns=[dullname])

    return ans

def addem(weights):
    ## Used for constraints
    return 1.0 - sum(weights)

def variance(weights, sigma):
    ## returns the variance (NOT standard deviation) given weights and sigma
    return (np.matrix(weights)*sigma*np.matrix(weights).transpose())[0,0]

def neg_SR(weights, sigma, mus):
    ## Returns minus the Sharpe Ratio (as we're minimising)

    """
    estreturn=250.0*((np.matrix(x)*mus)[0,0])
    variance=(variance(x,sigma)**.5)*16.0
    """
    estreturn=(np.matrix(weights)*mus)[0,0]
    std_dev=(variance(weights,sigma)**.5)

    return -estreturn/std_dev

def sigma_from_corr(std, corr):
    sigma=std*corr*std

    return sigma

def basic_opt(std,corr,mus):
    number_assets=mus.shape[0]
    sigma=sigma_from_corr(std, corr)
    start_weights=[1.0/number_assets]*number_assets

    ## Constraints - positive weights, adding to 1.0
    bounds=[(0.0,1.0)]*number_assets
    cdict=[{'type':'eq', 'fun':addem}]

    return minimize(neg_SR_riskfree, start_weights, (sigma, mus), method='SLSQP', bounds=bounds, constraints=cdict, tol=0.00001)

def neg_SR_riskfree(weights, sigma, mus, riskfree=0.005):
    ## Returns minus the Sharpe Ratio (as we're minimising)

    """
    estreturn=250.0*((np.matrix(x)*mus)[0,0])
    variance=(variance(x,sigma)**.5)*16.0
    """
    estreturn=(np.matrix(weights)*mus)[0,0] - riskfree
    std_dev=(variance(weights,sigma)**.5)

    return -estreturn/std_dev

def equalise_vols(returns, default_vol):
    """
    Normalises returns so they have the in sample vol of defaul_vol (annualised)
    Assumes daily returns
    """

    factors=(default_vol/16.0)/returns.std(axis=0)
    facmat=create_dull_pd_matrix(dullvalue=factors, dullname=returns.columns, index=returns.index)
    norm_returns=returns*facmat
    norm_returns.columns=returns.columns

    return norm_returns

def offdiag_matrix(offvalue, nlength):
    identity=np.diag([1.0]*nlength)
    for x in range(nlength):
        for y in range(nlength):
            if x!=y:
                identity[x][y]=offvalue
    return identity

def get_avg_corr(sigma):
    new_sigma=copy(sigma)
    np.fill_diagonal(new_sigma,np.nan)
    return np.nanmean(new_sigma)

def nearest_to_listvals(x, lvalues=[0.0, 0.25, 0.5, 0.75, 0.9]):
    ## return x rounded to nearest of lvalues

    if len(lvalues)==1:
        return lvalues[0]

    d1=abs(x - lvalues[0])
    d2=abs(x - lvalues[1])

    if d1<d2:
        return lvalues[0]

    newlvalues=lvalues[1:]

    return nearest_to_listvals(x, newlvalues)

def handcrafted(returns, equalisevols=True, default_vol=0.2):
    """
    Handcrafted optimiser
    """

    count_assets=len(returns.columns)

    try:

        assert equalisevols is True
        assert count_assets<=3
    except:
        raise Exception("Handcrafting only works with equalised vols and 3 or fewer assets")

    if count_assets<3:
        ## Equal weights
        return [1.0/count_assets]*count_assets

    est_corr=returns.corr().values
    c1=nearest_to_listvals(est_corr[0][1])
    c2=nearest_to_listvals(est_corr[0][2])
    c3=nearest_to_listvals(est_corr[1][2])

    wts_to_use=HANDCRAFTED_WTS[(HANDCRAFTED_WTS.c1==c1) & (HANDCRAFTED_WTS.c2==c2) & (HANDCRAFTED_WTS.c3==c3)].irow(0)

    return [wts_to_use.w1, wts_to_use.w2, wts_to_use.w3]

def opt_shrinkage(returns, shrinkage_factors, equalisevols=True, default_vol=0.2):
    """
    Returns the optimal portfolio for the dataframe returns using shrinkage

    shrinkage_factors is a tuple, shrinkage of mean and correlation

    If equalisevols=True then normalises returns to have same standard deviation; the weights returned
       will be 'risk weightings'

    """

    if equalisevols:
        use_returns=equalise_vols(returns, default_vol)
    else:
        use_returns=returns

    (shrinkage_mean, shrinkage_corr)=shrinkage_factors

    ## Sigma matrix
    ## Use correlation and then convert back to variance
    est_corr=use_returns.corr().values
    avg_corr=get_avg_corr(est_corr)
    prior_corr=offdiag_matrix(avg_corr, est_corr.shape[0])

    sigma_corr=shrinkage_corr*prior_corr+(1-shrinkage_corr)*est_corr
    cov_vector=use_returns.std().values

    sigma=cov_vector*sigma_corr*cov_vector

    ## mus vector
    avg_return=np.mean(use_returns.mean())
    est_mus=np.array([use_returns[asset_name].mean() for asset_name in use_returns.columns], ndmin=2).transpose()
    prior_mus=np.array([avg_return for asset_name in use_returns.columns], ndmin=2).transpose()

    mus=shrinkage_mean*prior_mus+(1-shrinkage_mean)*est_mus

    ## Starting weights
    number_assets=use_returns.shape[1]
    start_weights=[1.0/number_assets]*number_assets

    ## Constraints - positive weights, adding to 1.0
    bounds=[(0.0,1.0)]*number_assets
    cdict=[{'type':'eq', 'fun':addem}]

    ans=minimize(neg_SR, start_weights, (sigma, mus), method='SLSQP', bounds=bounds, constraints=cdict, tol=0.00001)

    return ans['x']

def handcraft_equal(returns):
    """
    dynamic handcrafting, equal weights only
    """

    ## RETURNS Correlation matrix
    use_returns=equalise_vols(returns, default_vol=16.0)

    ## Sigma matrix = correlations
    sigma=use_returns.cov()
    sigma[sigma<0.0]=0.0

    ungroupedreturns=dict([(x,returns[x]) for x in returns.columns])
    tree_data=hc_sigma(sigma, ungroupedreturns)

    tree_data=grouping_tree(tree_data)
    weights=tree_to_weights(tree_data)

    return weights

def hc_sigma(ungrouped_sigma, ungroupedreturns,  groupdata=None):
    """
    handcraft weights from sigma matrix

    Algo:
    - Find pair of assets with highest correlation
    - Form them into a new group with equal weights
    - The group becomes like a new asset
    - Once we only have two assets left, stop.

    Need to
    """

    if len(ungroupedreturns)==1:
        return groupdata[1]

    if groupdata is None:
        ## first run
        ## groupdata stores grouping information
        ## To begin with each group just consists of one asset
        groupdata=[[],list(ungrouped_sigma.columns)]

    groupedreturns=dict()

    ## iteration
    while len(ungroupedreturns)>0:
        ## current_sigma consists of the correlation of things we currently have
        if len(ungroupedreturns)==1:
            idx_list=[0]
        else:
            idx_list=find_highest_corr(ungrouped_sigma)

        name_list=tuple([ungrouped_sigma.columns[idx] for idx in idx_list])

        ## pair those things up
        (ungrouped_sigma, ungroupedreturns, groupedreturns,
         groupdata)=group_assets(ungrouped_sigma, ungroupedreturns, groupedreturns, groupdata, idx_list, name_list)

    new_returns=pd.concat(groupedreturns, axis=1)
    new_sigma=new_returns.corr()
    ## recursive
    return hc_sigma(new_sigma, groupedreturns,  groupdata=[[],groupdata[0]])

def find_highest_corr(sigmat):
    new_sigmat=copy(sigmat.values)
    np.fill_diagonal(new_sigmat, -100.0)
    (i,j)=np.unravel_index(new_sigmat.argmax(), new_sigmat.shape)
    return (i,j)

def group_assets(ungrouped_sigma, ungroupedreturns, groupedreturns, groupdata, idx_list, name_list):
    """
    Group assets
    """
    todelete=[]
    names=[]
    grouping=[]
    group_returns=[]

    weights=[1.0/len(idx_list)]*len(idx_list) ## could have more complex thing here...

    for (itemweight,idx, iname) in zip(weights,idx_list, name_list):
        gi=groupdata[1][idx]
        grouping.append(gi)
        gri=ungroupedreturns.pop(iname)

        group_returns.append(gri*itemweight)
        names.append(gri.name)

        ungrouped_sigma=ungrouped_sigma.drop(iname, axis=0)
        ungrouped_sigma=ungrouped_sigma.drop(iname, axis=1)

        todelete.append(idx)

    groupdata[0].append(grouping)
    gr_returns=pd.concat(group_returns, axis=1)
    gr_returns=gr_returns.sum(axis=1)

    gr_returns.name="[%s]" % "+".join(names)

    print "Pairing %s" % ", ".join(names)

    groupedreturns[gr_returns.name]=gr_returns
    groupdata[1]=[element for eindex, element in enumerate(groupdata[1]) if eindex not in todelete]

    return (ungrouped_sigma, ungroupedreturns, groupedreturns,
         groupdata)

def grouping_tree(tree_data, sigma):
    """
    Group branches of 2 into larger if possible
    """
    pass

def corrs_in_group(group, sigma):
    asset_list=sum(group, [])
    littlesigma=sigma.loc[asset_list, asset_list]

def corr_from_leaf(leaf, sigma):
    return sigma[leaf[0]][leaf[1]]

def tree_to_weights(tree_data):
    """
    convert a tree into weights
    """
    pass

def markosolver(returns, equalisemeans=False, equalisevols=True, default_vol=0.2, default_SR=1.0):
    """
    Returns the optimal portfolio for the dataframe returns

    If equalisemeans=True then assumes all assets have same return if False uses the asset means

    If equalisevols=True then normalises returns to have same standard deviation; the weights returned
       will be 'risk weightings'

    Note if usemeans=True and equalisevols=True effectively assumes all assets have same sharpe ratio

    """

    if equalisevols:
        use_returns=equalise_vols(returns, default_vol)
    else:
        use_returns=returns

    ## Sigma matrix
    sigma=use_returns.cov().values

    ## Expected mean returns
    est_mus=[use_returns[asset_name].mean() for asset_name in use_returns.columns]
    missingvals=[np.isnan(x) for x in est_mus]

    if equalisemeans:
        ## Don't use the data - Set to the average Sharpe Ratio
        mus=[default_vol*default_SR]*returns.shape[1]

    else:
        mus=est_mus

    mus=np.array(mus, ndmin=2).transpose()

    ## Starting weights
    number_assets=use_returns.shape[1]
    start_weights=[1.0/number_assets]*number_assets

    ## Constraints - positive weights, adding to 1.0
    bounds=[(0.0,1.0)]*number_assets
    cdict=[{'type':'eq', 'fun':addem}]

    ans=minimize(neg_SR, start_weights, (sigma, mus), method='SLSQP', bounds=bounds, constraints=cdict, tol=0.00001)
    wts=ans['x']

    return wts

def bootstrap_portfolio(returns_to_bs, monte_carlo=200, monte_length=250, equalisemeans=False, equalisevols=True, default_vol=0.2, default_SR=1.0):

    """
    Given dataframe of returns; returns_to_bs, performs a bootstrap optimisation

    We run monte_carlo numbers of bootstraps
    Each one contains monte_length days drawn randomly, with replacement
    (so *not* block bootstrapping)

    The other arguments are passed to the optimisation function markosolver

    Note - doesn't deal gracefully with missing data. Will end up downweighting stuff depending on how
      much data is missing in each boostrap. You'll need to think about how to solve this problem.

    """

    weightlist=[]
    for unused_index in range(monte_carlo):
        bs_idx=[int(random.uniform(0,1)*len(returns_to_bs)) for i in range(monte_length)]

        returns=returns_to_bs.iloc[bs_idx,:]
        weight=markosolver(returns, equalisemeans=equalisemeans, equalisevols=equalisevols, default_vol=default_vol, default_SR=default_SR)
        weightlist.append(weight)

    ### We can take an average here; only because our weights always add up to 1. If that isn't true
    ###    then you will need to some kind of renormalisation

    theweights_mean=list(np.mean(weightlist, axis=0))
    return theweights_mean

def optimise_over_periods(data, date_method, fit_method, rollyears=20, equalisemeans=False, equalisevols=True,
                          monte_carlo=100, monte_length=None, shrinkage_factors=(0.5, 0.5),
                          weightdf=None):
    """
    Do an optimisation

    Returns data frame of weights

    Note if fitting in sample weights will be somewhat boring

    Doesn't deal with eg missing data in certain subperiods

    """
    if monte_length is None:
        monte_length=int(len(data.index)*.1)

    ## Get the periods
    fit_periods=generate_fitting_dates(data, date_method, rollyears=rollyears)

    ## Do the fitting
    ## Build up a list of weights, which we'll concat
    weight_list=[]
    for fit_tuple in fit_periods:
        ## Fit on the slice defined by first two parts of the tuple
        period_subset_data=data[fit_tuple[0]:fit_tuple[1]]

        ## Can be slow, if bootstrapping, so indicate where we are

        print "Fitting data for %s to %s" % (str(fit_tuple[2]), str(fit_tuple[3]))

        if fit_method=="one_period":
            weights=markosolver(period_subset_data, equalisemeans=equalisemeans, equalisevols=equalisevols)
        elif fit_method=="bootstrap":
            weights=bootstrap_portfolio(period_subset_data, equalisemeans=equalisemeans,
                                        equalisevols=equalisevols, monte_carlo=monte_carlo,
                                        monte_length=monte_length)
        elif fit_method=="shrinkage":
            weights=opt_shrinkage(period_subset_data, shrinkage_factors=shrinkage_factors, equalisevols=equalisevols)

        elif fit_method=="fixed":
            weights=[float(weightdf[weightdf.Country==ticker].Weight.values) for ticker in list(period_subset_data.columns)]
        else:
            raise Exception("Fitting method %s unknown" % fit_method)

        ## We adjust dates slightly to ensure no overlaps
        dindex=[fit_tuple[2]+datetime.timedelta(seconds=1), fit_tuple[3]-datetime.timedelta(seconds=1)]

        ## create a double row to delineate start and end of test period
        weight_row=pd.DataFrame([weights]*2, index=dindex, columns=data.columns)

        weight_list.append(weight_row)

    weight_df=pd.concat(weight_list, axis=0)

    return weight_df

"""
Now we need to do this with expanding or rolling window
"""

"""
Generate the date tuples
"""

def generate_fitting_dates(data, date_method, rollyears=20):
    """
    generate a list 4 tuples, one element for each year in the data
    each tuple contains [fit_start, fit_end, period_start, period_end] datetime objects
    the last period will be a 'stub' if we haven't got an exact number of years

    date_method can be one of 'in_sample', 'expanding', 'rolling'

    if 'rolling' then use rollyears variable
    """

    start_date=data.index[0]
    end_date=data.index[-1]

    ## generate list of dates, one year apart, including the final date
    yearstarts=list(pd.date_range(start_date, end_date, freq="12M"))+[end_date]

    ## loop through each period
    periods=[]
    for tidx in range(len(yearstarts))[1:-1]:
        ## these are the dates we test in
        period_start=yearstarts[tidx]
        period_end=yearstarts[tidx+1]

        ## now generate the dates we use to fit
        if date_method=="in_sample":
            fit_start=start_date
        elif date_method=="expanding":
            fit_start=start_date
        elif date_method=="rolling":
            yearidx_to_use=max(0, tidx-rollyears)
            fit_start=yearstarts[yearidx_to_use]
        else:
            raise Exception("don't recognise date_method %s" % date_method)

        if date_method=="in_sample":
            fit_end=end_date
        elif date_method in ['rolling', 'expanding']:
            fit_end=period_start
        else:
            raise Exception("don't recognise date_method %s " % date_method)

        periods.append([fit_start, fit_end, period_start, period_end])

    ## give the user back the list of periods
    return periods

rawdata=read_ts_csv("/home/rob/workspace/systematictradingexamples/plots_for_perhaps/MSCI_data.csv")
refdata=pd.read_csv("/home/rob/workspace/systematictradingexamples/plots_for_perhaps/MSCI_ref.csv")
tickers=list(refdata[(refdata.EmorDEV=="DEV") & (refdata.Type=="Country")].Country.values) #mom 12bp
#tickers=list(refdata[refdata.Type=="Country"].Country.values) #mom 12bp

fix_hcweights=pd.read_csv("/home/rob/workspace/systematictradingexamples/plots_for_perhaps/devhcweights.csv")
fix_capweights=pd.read_csv("/home/rob/workspace/systematictradingexamples/plots_for_perhaps/devcapweights.csv")
fix_eqweights=pd.DataFrame(dict(Country=tickers, Weight=[1.0/len(tickers)]*len(tickers)))

data=calc_asset_returns(rawdata, tickers)

### IDEA: to boostrap the results
### Repeatedly draw from 'data' to make new pseudo series

oneperiodweights=optimise_over_periods(data, "expanding", "one_period", equalisemeans=False, equalisevols=True)
#bootstrapweights=optimise_over_periods(data, "expanding", "bootstrap", equalisemeans=True, equalisevols=True)
exposthcweights=optimise_over_periods(data, "expanding", "fixed", weightdf=fix_hcweights, equalisemeans=True, equalisevols=True)
equalweights=optimise_over_periods(data, "expanding", "fixed", weightdf=fix_eqweights, equalisemeans=True, equalisevols=True)
marketcapweights=optimise_over_periods(data, "expanding", "fixed", weightdf=fix_capweights, equalisemeans=True, equalisevols=True)

index_returns=(1.0+data).cumprod().ffill()
last_return=index_returns.irow(-1).values
last_return=pd.DataFrame(np.array([last_return]*len(data)), data.index)
last_return.columns=data.columns
index_returns = index_returns / last_return

marketcapweights = marketcapweights.reindex(index_returns.index, method="ffill")
marketcapweights=marketcapweights*index_returns

marketcapweights=marketcapweights.ffill()

## portfolio, take out missing weights
p1=portfolio_return(data, oneperiodweights)[pd.datetime(1994,1,1):]
#p2=portfolio_return(data, bootstrapweights)
p3=portfolio_return(data, exposthcweights)[pd.datetime(1994,1,1):]
p4=portfolio_return(data, equalweights)[pd.datetime(1994,1,1):]
p5=portfolio_return(data, marketcapweights)[pd.datetime(1994,1,1):]

drag1=p3 - p1
drag2=p4 - p5

def stats(x):
    ann_mean=x.mean()*12
    ann_std = x.std()*(12**.5)
    geo_mean = ann_mean - (ann_std**2)/2.0
    sharpe = geo_mean / ann_std

    return (ann_mean, ann_std, geo_mean, sharpe)

print stats(p1)
print stats(p3)
print stats(p4)
print stats(p5)

toplot=pd.concat([p1, p3, p4, p5], axis=1)
toplot.columns=["Optimised", "Handcraft", "Equal", "Market Cap"]
toplot.cumsum().plot()

show()

p1.cumsum().plot(color="black", ls="solid")
p3.cumsum().plot(color="gray", ls="solid")
p4.cumsum().plot(color="black", ls="dashed")
p5.cumsum().plot(color="gray", ls="dashed")

legend( ["Optimised", "Handcraft", "Equal", "Market Cap"], loc="upper left")
frame=plt.gca()

#frame.get_yaxis().set_visible(False)

rcParams.update({'font.size': 18})

file_process("compareoptmethods")
show()

drag1.cumsum().plot(color="gray", ls="solid")

legend( [ "Handcraft vs MktCap"], loc="upper left")
frame=plt.gca()

#frame.get_yaxis().set_visible(False)

rcParams.update({'font.size': 18})

file_process("compareoptmethodstracking")
show()