#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Short interest dynamics

"""
from __future__ import print_function, division

import os
import zipfile

import datetime as dt
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

path = os.getenv("HOME") + '/Dropbox/Research/data/Compustat/data/'
# __location__ = os.path.realpath(os.path.join(os.getcwd(),
#                                 os.path.dirname(__file__)))
# path = os.path.join(__location__, path + 'Compustat/data/')


def date_convert(string):
    return dt.datetime.strptime(string, '%d-%m-%Y')


def import_data():
    """Import data and save it to the disk.

    """
    zf = zipfile.ZipFile(path + 'short_int.zip', 'r')
    name = zf.namelist()[0]

    short_int = pd.read_csv(zf.open(name),
                            converters={'datadate': date_convert})
    columns = {'datadate': 'date',
               'SHORTINTADJ': 'short_int',
               'GVKEY': 'gvkey'}
    short_int.rename(columns=columns, inplace=True)
    short_int.set_index(['gvkey', 'date'], inplace=True)
    short_int.sort_index(inplace=True)

    short_int.to_hdf(path + 'short_int.h5', key='short_int')

    print(short_int.head())
    print(short_int.dtypes)
    print('Number of unique companies: ',
          short_int.index.get_level_values('gvkey').nunique())
    print('Number of unique dates: ',
          short_int.index.get_level_values('date').nunique())
    print('Min and Max date: ',
          short_int.index.get_level_values('date').min().date(), ',',
          short_int.index.get_level_values('date').max().date())


def load_data():
    """Load data from disk and check for sanity.

    """
    return pd.read_hdf(path + 'short_int.h5', 'short_int')


def count_companies(short_int):
    """Plot number of companies over time.

    """
    df = short_int.reset_index().groupby('date')['gvkey'].nunique()

    sns.set_context('paper')

    df.plot(figsize=(10, 3))
    plt.show()

    data = df.ix[dt.date(2006, 1, 1):dt.date(2007, 6, 30)]
    data.plot(figsize=(10, 3))
    plt.show()


def mean_short_int(short_int):
    """Mean short interest on each date.

    """
    df = short_int.groupby(level='date')['short_int'].mean()

    sns.set_context('paper')

    df.plot(figsize=(10, 3))
    plt.show()

    df.ix[:dt.date(2004, 12, 31)].plot(figsize=(10, 3))
    plt.show()

    df.ix[dt.date(2006, 1, 1):dt.date(2007, 6, 30)].plot(figsize=(10, 3))
    plt.show()


if __name__ == '__main__':

    import_data()
    short_int = load_data()
    count_companies(short_int)
    mean_short_int(short_int)
