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

# ohlc_chart_maker.py
# ==========================
# * Version     : 0.2
# * Last update : 2019-01-26
#

import numpy
import mpl_finance

from matplotlib import dates
from matplotlib import pyplot

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

def make_mpl_dates(year, month, day):
    """ @return Matplotlib dates """
    return dates.date2num(DateTime(year, month, day))

def get_test_data():
    """ @return test 2d array (Time, Open, High, Low, Close, Volume) """
    data = numpy.array([
        [make_mpl_dates(2019, 1, 14),  90, 105,  85, 100, 1000],
        [make_mpl_dates(2019, 1, 15), 110, 110, 110, 110, 1100],
        [make_mpl_dates(2019, 1, 16), 125, 125, 115, 120, 1200],
        [make_mpl_dates(2019, 1, 17), 120, 140, 120, 140, 1300],
        [make_mpl_dates(2019, 1, 18), 130, 145, 130, 135, 1200],
        [make_mpl_dates(2019, 1, 21), 125, 125, 110, 110, 1100],
        [make_mpl_dates(2019, 1, 22), 115, 120,  90, 105,  900],
        [make_mpl_dates(2019, 1, 23), 110, 115, 100, 110, 1100],
        [make_mpl_dates(2019, 1, 24), 130, 135, 125, 125, 1400],
        [make_mpl_dates(2019, 1, 25), 120, 140, 115, 135, 1500]])
    return data

if __name__ == "__main__":
    # Time, Open, High, Low, Close, Volume
    data = get_test_data()

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

    # x-axis settings
    ax1.set_xticklabels(data[:, 0], rotation=90, size="small")
    # ax1.xaxis.set_major_locator(DayLocator())
    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_ystep = 10

    price_max  = numpy.amax(data[:, 2])
    price_min  = numpy.amin(data[:, 3])
    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) + 1) * 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=8,
                                      colorup="hotpink",
                                      colordown="cornflowerblue")
    ax1.grid(color="silver", linestyle=":")

    # ax2 settings
    volume_ystep = 500

    volume_max  =  numpy.amax(data[:, 5])
    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.25,
            color="forestgreen")
    ax2.grid(axis="y", color="orange", linestyle="--")

    pyplot.show()