#!/usr/bin/env python

# Copyright 2010, 2011 Luther Thompson

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import csv
import json
import math
import os
import os.path
import shutil
import sys
import time
import urllib.request

# Disallow 2-digit years for future compatability
time.accept2dyear = False
# Store the present time to a constant to make sure it's the same everywhere
NOW = time.time()

USER_DATA_DIR = '{0}/.btctrade'.format(os.environ['HOME'])


class BTCPrice:
  '''Stores data for an individual trade'''

  def __init__(self, price, time):
    '''Arguments are the price and time of the trade'''
    self.price = price
    self._time = time

  def __str__(self):
    return '{:<11.6g} {}'.format(self.price, time.strftime(
        '%Y-%_m-%e %a %k:%M', time.localtime(self._time)))


class RunningTotal:
  '''Stores a running total for all trades of a currency.

  Can report high, low, and average.

  '''

  def __init__(self):
    '''Starts off empty'''
    self.low = None
    self.high = None
    self.sum = 0
    self.volume = 0

  def add(self, price, time, volume):
    '''Processes a new trade.

    price: The price in bitcoins
    time: Time of the trade
    volume: Amount of bitcoins in the trade

    '''
    priceData = BTCPrice(price, time)

    if not self.low or priceData.price < self.low.price:
      self.low = priceData
    if not self.high or priceData.price > self.high.price:
      self.high = priceData

    # Used for calculating the average. As written, it gives a higher weight to
    # when BTC are low. If I convert the quantity from BTC to the commodity, it
    # gives a higher weight to when the commodity is low.
    self.sum += priceData.price * volume
    self.volume += volume

  def __str__(self):
    if not self.volume:
      return 'no volume'

    average = self.sum / self.volume
    uncertainty = self.high.price - self.low.price
    precision = (max(
        math.ceil(math.log10(average)) - math.floor(math.log10(uncertainty)), 1)
                 if uncertainty else 6)

    return '''\
High:    {0.high}
Low:     {0.low}
Average: {1:.{2}g}
'''.format(self, average, precision)


class RunningTotalDict(dict):
  '''A dict with an infinite supply of RunningTotal objects.

  I assume this is faster than calling dict.setdefault many times with the same
  key.

  '''
  def __missing__(self, key):
    self[key] = RunningTotal()
    return self[key]


def comma_list(s):
  return s.split(',')


def download(url, min_timestamp, offline, filename=None):
  '''Downloads data from http://bitcoincharts.com/t/.

  Stores data in USER_DATA_DIR.
  url: Specifies the data to download. See
  http://bitcoincharts.com/about/markets-api/ for details.
  min_timestamp: If the file was already downloaded after this time, don't
  download it now. If this is later than 15 min ago, it will be reset to that
  time.
  offline: If True, don't do anything.
  filename: The filename to store the data in.

  '''
  if not offline:
    filename = filename or url
    min_timestamp = min(min_timestamp, NOW - 900)
    filepath = '{}/{}'.format(USER_DATA_DIR, filename)
    if (not os.path.exists(filepath) or
        os.path.getmtime(filepath) < min_timestamp):
      print('Downloading {}...'.format(url), end='')
      sys.stdout.flush()
      try:
        data = urllib.request.urlopen(
          'http://bitcoincharts.com/t/{}'.format(url))
      except urllib.error.URLError as e:
        print(e)
        return
      data_file = open(filepath, 'wb')
      shutil.copyfileobj(data, data_file)
      data_file.close()
      print('done')


## Options ##
ap = argparse.ArgumentParser()
ap.add_argument(
  '-V', '--version', action='version', version='%(prog)s 0.4.0')
ap.add_argument('-c', '--currency', type=comma_list)
ap.add_argument('-m', '--market', type=comma_list)
ap.add_argument('-o', '--offline', action='store_true')
ap.add_argument('-s', '--separate', action='store_true')
ap.add_argument('time', type=int, nargs='*')
arg = ap.parse_args()

# Decide the start and finish times
if arg.time:
  TIME_LENGTH = 6
  given_length = len(arg.time)

  if given_length > TIME_LENGTH:
    print('Error: Argument format is <year month day hour minute second>. Smaller units may be omitted for longer periods.')
    exit(2)

  # End of the period
  finish_struct = arg.time[:]
  finish_struct[-1] += 1

  # Prepare for mktime()
  for struct in [arg.time, finish_struct]:
    # Complete the struct_time by providing the lowest possible values for the
    # missing fields.
    for i in range(given_length, TIME_LENGTH):
      struct.append((None, 1, 1, 0, 0, 0)[i])
    struct.extend([0, 0, -1])

  start = time.mktime(tuple(arg.time))
  finish = time.mktime(tuple(finish_struct))
  if start > NOW:
    print('warning: Period is in the future')
else:
  # No time given. Use entire history.
  start = 0
  finish = NOW

totals = RunningTotalDict()

if not os.path.isdir(USER_DATA_DIR):
  os.mkdir(USER_DATA_DIR)

# We need to find out what markets are available and when their latest trades
# were.
download('markets.json', finish, arg.offline)
market_file = open('{}/markets.json'.format(USER_DATA_DIR))
market_data = json.load(market_file)
market_file.close()

for market in market_data:
  currency = market['currency']
  symbol = market['symbol']
  if ((not arg.currency or currency in arg.currency) and
      (not arg.market or symbol in arg.market)):
    url = 'trades.csv?start=0&symbol={}'.format(symbol)
    filename = '{}.csv'.format(symbol)
    latest = market['latest_trade']

    download(url, min(latest, finish), arg.offline, filename)

    data_file = open('{}/{}'.format(USER_DATA_DIR, filename))
    data = csv.DictReader(
      data_file, ('time', 'price', 'amount'), quoting=csv.QUOTE_NONNUMERIC)

    for row in data:
      # Assumption: chronological order
      trade_time = row['time']

      if trade_time >= finish:
        break
      elif trade_time >= start:
        key = symbol if arg.separate else currency
        totals[key].add(1 / row['price'], trade_time, row['amount'])

    data_file.close()

currencies = list(totals.keys())
currencies.sort()
for currency in currencies:
  print(currency)
  print(totals[currency])
