# -*- coding: utf-8 -*-

# ohlc_chart_maker.py
# ==========================
# * Version     : 0.3
# * Last update : 2019-03-01
#

import csv
import mpl_finance
import numpy
import sys

from matplotlib import dates
from matplotlib import pyplot

from datetime import datetime as DateTime
from matplotlib.dates import DateFormatter
from matplotlib.dates import WeekdayLocator

def read_csv_data(fname):
    """ @return csv data 2d array (Time, Open, High, Low, Close, Volume) """
    csv_data = []

    with open(fname, "r") as reader:
        csv_reader = csv.reader(reader)
        for row in csv_reader:
            list_temp = []
            list_temp.append(
                dates.date2num(DateTime.strptime(row[0], "%Y-%m-%d")))
            list_temp.append(int(row[1]))
            list_temp.append(int(row[2]))
            list_temp.append(int(row[3]))
            list_temp.append(int(row[4]))
            list_temp.append(int(row[5]))

            csv_data.append(list_temp)

    return numpy.array(csv_data)

if __name__ == "__main__":
    argv = sys.argv
    argc = len(argv)

    csv_file_name = ""
    price_ystep = 0
    volume_ystep = 0

    if argc < 2:
        print("Usage: python3",
              argv[0],
              "[historical data file] ([price y step]) ([volume y step])")
        print("e.g.1: python3", argv[0], "test_data.csv")
        print("e.g.2: python3", argv[0], "test_data.csv 40 10000")
        sys.exit(0)
    else:
        csv_file_name = argv[1]

    if argc > 2:
        price_ystep = int(argv[2])

    if argc > 3:
        volume_ystep = int(argv[3])

    # Time, Open, High, Low, Close, Volume
    data = read_csv_data(csv_file_name)

    # pyplot settings
    pyplot.rcParams["font.family"] = "monospace"
    pyplot.rcParams["axes.xmargin"] = 0.01

    fig = pyplot.figure(figsize=(14, 4))
    ax1 = fig.add_subplot(1, 1, 1)
    ax2 = ax1.twinx()

    # x-axis settings
    ax1.tick_params("x", labelsize="small")
    ax1.xaxis.set_major_locator(WeekdayLocator(byweekday=dates.MO))
    ax1.xaxis.set_major_formatter(DateFormatter("%b-%d"))

    # 「OHLC チャート」と「出来高」の表示比 (2 なら 2:1 で表示する)
    pv_yratio = 2

    # ax1 settings
    price_max = numpy.amax(data[:, 2])
    price_min = numpy.amin(data[:, 3])

    if price_ystep == 0:
        price_ystep = int((price_max - price_min) / 6)

    price_ymax = (int(price_max/price_ystep) + 1) * price_ystep
    price_ymin = price_min - (price_ymax - price_min) / pv_yratio
    price_ymin = int(price_ymin / price_ystep) * price_ystep
    ax1.set_ylim(price_ymin, price_ymax)

    price_yticks = []
    price_ytick = (int(price_min/price_ystep) + 0) * price_ystep
    while price_ytick <= price_ymax:
        price_yticks.append(price_ytick)
        price_ytick += price_ystep
    ax1.set_yticks(price_yticks)

    ax1.set_axisbelow(True)
    mpl_finance.plot_day_summary_ohlc(ax1,
                                      data,
                                      ticksize=3.25,
                                      colorup="hotpink",
                                      colordown="cornflowerblue")
    ax1.grid(color="silver", linestyle=":")

    # ax2 settings
    volume_max = numpy.amax(data[:, 5])

    if volume_ystep == 0:
        volume_ystep = int(volume_max / 3)

    volume_ymax = int((volume_max + volume_ystep - 1) / volume_ystep)
    volume_ymax *= volume_ystep * (pv_yratio + 1)
    ax2.set_ylim(0, volume_ymax)

    volume_yticks = []
    volume_ytick = 0
    volume_ytickmax = int((volume_max + volume_ystep - 1) / volume_ystep)
    volume_ytickmax *= volume_ystep
    while volume_ytick <= volume_ytickmax:
        volume_yticks.append(volume_ytick)
        volume_ytick += volume_ystep
    ax2.set_yticks(volume_yticks)

    ax2.set_axisbelow(True)
    ax2.bar(data[:, 0],
            data[:, 5],
            width=0,
            edgecolor="forestgreen")

    ax2.grid(axis="y", color="orange", linestyle="--")

    fig.autofmt_xdate(rotation=90, ha="center")
    fig.tight_layout()
    #pyplot.show()

    png_file_name = "fig_" + csv_file_name.replace(".csv", ".png")
    pyplot.savefig(png_file_name)