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

# ohlc_chart_maker.py
# ==========================
# * Version     : 0.2-test1
# * Last update : 2019-03-03
#

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 WeekdayLocator
from matplotlib.ticker import MaxNLocator

JP_FONT_FAMILY = "Ume Gothic"

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, 1400],
        [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

def test_balloon(ax, time, price, price_type):
    offset_y = 16
    if price_type[0] == "l":
        offset_y *= -1

    color = "gray"
    bbox_props = dict(boxstyle="square",
                      color=color,
                      linewidth=1.5)
    arrow_props = dict(arrowstyle="wedge, tail_width=1.0",
                       color=color,
                       patchA=None,
                       linewidth=1.5)
    annotate_args = dict(s=price,
                         xy=(time, price),
                         xycoords="data",
                         xytext=(0, offset_y),
                         textcoords="offset points",
                         size=8,
                         ha="center",
                         va="center",
                         bbox=bbox_props,
                         arrowprops=arrow_props)
    ax.annotate(**annotate_args)

    color = "white"
    bbox_props["color"] = color
    bbox_props["linewidth"] = 0
    arrow_props["color"] = color
    arrow_props["linewidth"] = 0
    ax.annotate(**annotate_args)

def test_annotation(ax, time, price, text, offset_x, offset_y):
    bbox_props = dict(boxstyle="round",
                      facecolor="white")
    arrow_props = dict(arrowstyle="-",
                       shrinkA=0,
                       connectionstyle="angle, angleA=0, angleB=90")
    ax.annotate(text,
                xy=(time, price),
                xycoords="data",
                xytext=(offset_x, offset_y),
                textcoords="offset points",
                fontfamily=JP_FONT_FAMILY,
                bbox=bbox_props,
                arrowprops=arrow_props)

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()

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

    # ax1 settings
    mpl_finance.plot_day_summary_ohlc(ax1,
                                      data,
                                      ticksize=8,
                                      colorup="hotpink",
                                      colordown="cornflowerblue")

    (bottom, top) = ax1.get_ylim()
    new_bottom = bottom - (top - bottom) / pv_yratio
    ax1.set_ylim(bottom=new_bottom)
    ax1.yaxis.set_major_locator(MaxNLocator(13))
    ticks = ax1.get_yticks()
    indices = numpy.where((ticks >= bottom) & (ticks <= top))
    ax1.set_yticks(ticks[indices])

    ax1.set_axisbelow(True)
    ax1.grid(color="silver", linestyle=":")

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

    # ax2 settings
    ax2.bar(data[:, 0],
            data[:, 5],
            width=0.25,
            color="forestgreen")

    (bottom, top) = ax2.get_ylim()
    new_top = top * (pv_yratio + 1)
    ax2.set_ylim(top=new_top)
    ax2.yaxis.set_major_locator(MaxNLocator(13))
    ticks = ax2.get_yticks()
    indices = numpy.where(ticks <= top)
    ax2.set_yticks(ticks[indices])

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

    # test 1
    test_annotation(ax1,
                    make_mpl_dates(2019, 1, 18),
                    145,
                    u"1/18 1Q 決算発表",
                    10,
                    30)
    test_annotation(ax1,
                    make_mpl_dates(2019, 1, 21),
                    110,
                    u"1/21 ？？？のお知らせ",
                    10,
                    -90)
    # test 2
    test_balloon(ax1, make_mpl_dates(2019, 1, 18), 145, "high")
    test_balloon(ax1, make_mpl_dates(2019, 1, 14), 85, "low")

    pyplot.tight_layout()
    pyplot.show()