text: #!/usr/bin/env python
  # -*- coding: utf-8; py-indent-offset:4 -*-
  ###############################################################################
  #
  # Copyright (C) 2015-2020 Daniel Rodriguez
  #
  # 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/>.
  #
  ###############################################################################
  from __future__ import (absolute_import, division, print_function,
                          unicode_literals)
  
  from ..utils.py3 import range, zip
  
  import matplotlib.collections as mcol
  import matplotlib.colors as mcolors
  import matplotlib.legend as mlegend
  import matplotlib.lines as mlines
  
  from .utils import shade_color
  
  
  class CandlestickPlotHandler(object):
      legend_opens = [0.50, 0.50, 0.50]
      legend_highs = [1.00, 1.00, 1.00]
      legend_lows = [0.00, 0.00, 0.00]
      legend_closes = [0.80, 0.00, 1.00]
  
      def __init__(self,
                   ax, x, opens, highs, lows, closes,
                   colorup='k', colordown='r',
                   edgeup=None, edgedown=None,
                   tickup=None, tickdown=None,
                   width=1, tickwidth=1,
                   edgeadjust=0.05, edgeshading=-10,
                   alpha=1.0,
                   label='_nolegend',
                   fillup=True,
                   filldown=True,
                   **kwargs):
  
          # Manager up/down bar colors
          r, g, b = mcolors.colorConverter.to_rgb(colorup)
          self.colorup = r, g, b, alpha
          r, g, b = mcolors.colorConverter.to_rgb(colordown)
          self.colordown = r, g, b, alpha
          # Manage the edge up/down colors for the bars
          if edgeup:
              r, g, b = mcolors.colorConverter.to_rgb(edgeup)
              self.edgeup = ((r, g, b, alpha),)
          else:
              self.edgeup = shade_color(self.colorup, edgeshading)
  
          if edgedown:
              r, g, b = mcolors.colorConverter.to_rgb(edgedown)
              self.edgedown = ((r, g, b, alpha),)
          else:
              self.edgedown = shade_color(self.colordown, edgeshading)
  
              # Manage the up/down tick colors
          if tickup:
              r, g, b = mcolors.colorConverter.to_rgb(tickup)
              self.tickup = ((r, g, b, alpha),)
          else:
              self.tickup = self.edgeup
  
          if tickdown:
              r, g, b = mcolors.colorConverter.to_rgb(tickdown)
              self.tickdown = ((r, g, b, alpha),)
          else:
              self.tickdown = self.edgedown
  
          self.barcol, self.tickcol = self.barcollection(
              x, opens, highs, lows, closes,
              width, tickwidth, edgeadjust,
              label=label,
              fillup=fillup, filldown=filldown,
              **kwargs)
  
          # add collections to the axis and return them
          ax.add_collection(self.tickcol)
          ax.add_collection(self.barcol)
  
          # Update the axis
          ax.update_datalim(((0, min(lows)), (len(opens), max(highs))))
          ax.autoscale_view()
  
          # Add self as legend handler for this object
          mlegend.Legend.update_default_handler_map({self.barcol: self})
  
      def legend_artist(self, legend, orig_handle, fontsize, handlebox):
          x0 = handlebox.xdescent
          y0 = handlebox.ydescent
          width = handlebox.width / len(self.legend_opens)
          height = handlebox.height
  
          # Generate the x axis coordinates (handlebox based)
          xs = [x0 + width * (i + 0.5) for i in range(len(self.legend_opens))]
  
          barcol, tickcol = self.barcollection(
              xs,
              self.legend_opens, self.legend_highs,
              self.legend_lows, self.legend_closes,
              width=width, tickwidth=2,
              scaling=height, bot=y0)
  
          barcol.set_transform(handlebox.get_transform())
          handlebox.add_artist(barcol)
          tickcol.set_transform(handlebox.get_transform())
          handlebox.add_artist(tickcol)
  
          return barcol, tickcol
  
      def barcollection(self,
                        xs,
                        opens, highs, lows, closes,
                        width, tickwidth=1, edgeadjust=0,
                        label='_nolegend',
                        scaling=1.0, bot=0,
                        fillup=True, filldown=True,
                        **kwargs):
  
          # Prepack different zips of the series values
          oc = lambda: zip(opens, closes)  # NOQA: E731
          xoc = lambda: zip(xs, opens, closes)  # NOQA: E731
          iohlc = lambda: zip(xs, opens, highs, lows, closes)  # NOQA: E731
  
          colorup = self.colorup if fillup else 'None'
          colordown = self.colordown if filldown else 'None'
          colord = {True: colorup, False: colordown}
          colors = [colord[o < c] for o, c in oc()]
  
          edgecolord = {True: self.edgeup, False: self.edgedown}
          edgecolors = [edgecolord[o < c] for o, c in oc()]
  
          tickcolord = {True: self.tickup, False: self.tickdown}
          tickcolors = [tickcolord[o < c] for o, c in oc()]
  
          delta = width / 2 - edgeadjust
  
          def barbox(i, open, close):
              # delta seen as closure
              left, right = i - delta, i + delta
              open = open * scaling + bot
              close = close * scaling + bot
              return (left, open), (left, close), (right, close), (right, open)
  
          barareas = [barbox(i, o, c) for i, o, c in xoc()]
  
          def tup(i, open, high, close):
              high = high * scaling + bot
              open = open * scaling + bot
              close = close * scaling + bot
  
              return (i, high), (i, max(open, close))
  
          tickrangesup = [tup(i, o, h, c) for i, o, h, l, c in iohlc()]
  
          def tdown(i, open, low, close):
              low = low * scaling + bot
              open = open * scaling + bot
              close = close * scaling + bot
  
              return (i, low), (i, min(open, close))
  
          tickrangesdown = [tdown(i, o, l, c) for i, o, h, l, c in iohlc()]
  
          # Extra variables for the collections
          useaa = 0,  # use tuple here
          lw = 0.5,   # and here
          tlw = tickwidth,
  
          # Bar collection for the candles
          barcol = mcol.PolyCollection(
              barareas,
              facecolors=colors,
              edgecolors=edgecolors,
              antialiaseds=useaa,
              linewidths=lw,
              label=label,
              **kwargs)
  
          # LineCollections have a higher zorder than PolyCollections
          # to ensure the edges of the bars are not overwriten by the Lines
          # we need to put the bars slightly over the LineCollections
          kwargs['zorder'] = barcol.get_zorder() * 0.9999
  
          # Up/down ticks from the body
          tickcol = mcol.LineCollection(
              tickrangesup + tickrangesdown,
              colors=tickcolors,
              linewidths=tlw,
              antialiaseds=useaa,
              **kwargs)
  
          # return barcol, tickcol
          return barcol, tickcol
  
  
  def plot_candlestick(ax,
                       x, opens, highs, lows, closes,
                       colorup='k', colordown='r',
                       edgeup=None, edgedown=None,
                       tickup=None, tickdown=None,
                       width=1, tickwidth=1.25,
                       edgeadjust=0.05, edgeshading=-10,
                       alpha=1.0,
                       label='_nolegend',
                       fillup=True,
                       filldown=True,
                       **kwargs):
  
      chandler = CandlestickPlotHandler(
          ax, x, opens, highs, lows, closes,
          colorup, colordown,
          edgeup, edgedown,
          tickup, tickdown,
          width, tickwidth,
          edgeadjust, edgeshading,
          alpha,
          label,
          fillup,
          filldown,
          **kwargs)
  
      # Return the collections. the barcol goes first because
      # is the larger,  has the dominant zorder and defines the legend
      return chandler.barcol, chandler.tickcol
  
  
  class VolumePlotHandler(object):
      legend_vols = [0.5, 1.0, 0.75]
      legend_opens = [0, 1, 0]
      legend_closes = [1, 0, 1]
  
      def __init__(self,
                   ax, x, opens, closes, volumes,
                   colorup='k', colordown='r',
                   edgeup=None, edgedown=None,
                   edgeshading=-5, edgeadjust=0.05,
                   width=1, alpha=1.0,
                   **kwargs):
  
          # Manage the up/down colors
          r, g, b = mcolors.colorConverter.to_rgb(colorup)
          self.colorup = r, g, b, alpha
          r, g, b = mcolors.colorConverter.to_rgb(colordown)
          self.colordown = r, g, b, alpha
  
          # Prepare the edge colors
          if not edgeup:
              self.edgeup = shade_color(self.colorup, edgeshading)
          else:
              r, g, b = mcolors.colorConverter.to_rgb(edgeup)
              self.edgeup = r, g, b, alpha
  
          if not edgedown:
              self.edgedown = shade_color(self.colordown, edgeshading)
          else:
              r, g, b = mcolors.colorConverter.to_rgb(edgedown)
              self.edgedown = r, g, b, alpha
  
          corners = (0, 0), (len(closes), max(volumes))
          ax.update_datalim(corners)
          ax.autoscale_view()
  
          self.barcol = self.barcollection(
              x, opens, closes, volumes,
              width=width, edgeadjust=edgeadjust,
              **kwargs)
  
          # add to axes
          ax.add_collection(self.barcol)
  
          # Add a legend handler for this object
          mlegend.Legend.update_default_handler_map({self.barcol: self})
  
      def legend_artist(self, legend, orig_handle, fontsize, handlebox):
          x0 = handlebox.xdescent
          y0 = handlebox.ydescent
          width = handlebox.width / len(self.legend_vols)
          height = handlebox.height
  
          # Generate the x axis coordinates (handlebox based)
          xs = [x0 + width * (i + 0.5) for i in range(len(self.legend_vols))]
  
          barcol = self.barcollection(
              xs, self.legend_opens, self.legend_closes, self.legend_vols,
              width=width, vscaling=height, vbot=y0)
  
          barcol.set_transform(handlebox.get_transform())
          handlebox.add_artist(barcol)
  
          return barcol
  
      def barcollection(self,
                        x, opens, closes, vols,
                        width, edgeadjust=0,
                        vscaling=1.0, vbot=0,
                        **kwargs):
  
          # Prepare the data
          openclose = lambda: zip(opens, closes)  # NOQA: E731
  
          # Calculate bars colors
          colord = {True: self.colorup, False: self.colordown}
          colors = [colord[open < close] for open, close in openclose()]
          edgecolord = {True: self.edgeup, False: self.edgedown}
          edgecolors = [edgecolord[open < close] for open, close in openclose()]
  
          # bar width to the sides
          delta = width / 2 - edgeadjust
  
          # small auxiliary func to return the bar coordinates
          def volbar(i, v):
              left, right = i - delta, i + delta
              v = vbot + v * vscaling
              return (left, vbot), (left, v), (right, v), (right, vbot)
  
          barareas = [volbar(i, v) for i, v in zip(x, vols)]
          barcol = mcol.PolyCollection(
              barareas,
              facecolors=colors,
              edgecolors=edgecolors,
              antialiaseds=(0,),
              linewidths=(0.5,),
              **kwargs)
  
          return barcol
  
  
  def plot_volume(
          ax, x, opens, closes, volumes,
          colorup='k', colordown='r',
          edgeup=None, edgedown=None,
          edgeshading=-5, edgeadjust=0.05,
          width=1, alpha=1.0,
          **kwargs):
  
      vhandler = VolumePlotHandler(
          ax, x, opens, closes, volumes,
          colorup, colordown,
          edgeup, edgedown,
          edgeshading, edgeadjust,
          width, alpha,
          **kwargs)
  
      return vhandler.barcol,
  
  
  class OHLCPlotHandler(object):
      legend_opens = [0.50, 0.50, 0.50]
      legend_highs = [1.00, 1.00, 1.00]
      legend_lows = [0.00, 0.00, 0.00]
      legend_closes = [0.80, 0.20, 0.90]
  
      def __init__(self,
                   ax, x, opens, highs, lows, closes,
                   colorup='k', colordown='r',
                   width=1, tickwidth=0.5,
                   alpha=1.0,
                   label='_nolegend',
                   **kwargs):
  
          # Manager up/down bar colors
          r, g, b = mcolors.colorConverter.to_rgb(colorup)
          self.colorup = r, g, b, alpha
          r, g, b = mcolors.colorConverter.to_rgb(colordown)
          self.colordown = r, g, b, alpha
  
          bcol, ocol, ccol = self.barcollection(
              x, opens, highs, lows, closes,
              width=width, tickwidth=tickwidth,
              label=label,
              **kwargs)
  
          self.barcol = bcol
          self.opencol = ocol
          self.closecol = ccol
  
          # add collections to the axis and return them
          ax.add_collection(self.barcol)
          ax.add_collection(self.opencol)
          ax.add_collection(self.closecol)
  
          # Update the axis
          ax.update_datalim(((0, min(lows)), (len(opens), max(highs))))
          ax.autoscale_view()
  
          # Add self as legend handler for this object
          mlegend.Legend.update_default_handler_map({self.barcol: self})
  
      def legend_artist(self, legend, orig_handle, fontsize, handlebox):
          x0 = handlebox.xdescent
          y0 = handlebox.ydescent
          width = handlebox.width / len(self.legend_opens)
          height = handlebox.height
  
          # Generate the x axis coordinates (handlebox based)
          xs = [x0 + width * (i + 0.5) for i in range(len(self.legend_opens))]
  
          barcol, opencol, closecol = self.barcollection(
              xs,
              self.legend_opens, self.legend_highs,
              self.legend_lows, self.legend_closes,
              width=1.5, tickwidth=2,
              scaling=height, bot=y0)
  
          barcol.set_transform(handlebox.get_transform())
          handlebox.add_artist(barcol)
          # opencol.set_transform(handlebox.get_transform())
          handlebox.add_artist(opencol)
          # closecol.set_transform(handlebox.get_transform())
          handlebox.add_artist(closecol)
  
          return barcol, opencol, closecol
  
      def barcollection(self,
                        xs,
                        opens, highs, lows, closes,
                        width, tickwidth,
                        label='_nolegend',
                        scaling=1.0, bot=0,
                        **kwargs):
  
          # Prepack different zips of the series values
          ihighlow = lambda: zip(xs, highs, lows)  # NOQA: E731
          iopen = lambda: zip(xs, opens)  # NOQA: E731
          iclose = lambda: zip(xs, closes)  # NOQA: E731
          openclose = lambda: zip(opens, closes)  # NOQA: E731
  
          colord = {True: self.colorup, False: self.colordown}
          colors = [colord[open < close] for open, close in openclose()]
  
          # Extra variables for the collections
          useaa = 0,
          lw = width,
          tlw = tickwidth,
  
          # Calculate the barranges
          def barrange(i, high, low):
              return (i, low * scaling + bot), (i, high * scaling + bot)
  
          barranges = [barrange(i, high, low) for i, high, low in ihighlow()]
  
          barcol = mcol.LineCollection(
              barranges,
              colors=colors,
              linewidths=lw,
              antialiaseds=useaa,
              label=label,
              **kwargs)
  
          def tickopen(i, open):
              open = open * scaling + bot
              return (i - tickwidth, open), (i, open)
  
          openticks = [tickopen(i, open) for i, open in iopen()]
          opencol = mcol.LineCollection(
              openticks,
              colors=colors,
              antialiaseds=useaa,
              linewidths=tlw,
              label='_nolegend',
              **kwargs)
  
          def tickclose(i, close):
              close = close * scaling + bot
              return (i, close), (i + tickwidth, close)
  
          closeticks = [tickclose(i, close) for i, close in iclose()]
          closecol = mcol.LineCollection(
              closeticks,
              colors=colors,
              antialiaseds=useaa,
              linewidths=tlw,
              label='_nolegend',
              **kwargs)
  
          # return barcol, tickcol
          return barcol, opencol, closecol
  
  
  def plot_ohlc(ax, x, opens, highs, lows, closes,
                colorup='k', colordown='r',
                width=1.5, tickwidth=0.5,
                alpha=1.0,
                label='_nolegend',
                **kwargs):
  
      handler = OHLCPlotHandler(
          ax, x, opens, highs, lows, closes,
          colorup, colordown,
          width, tickwidth,
          alpha,
          label,
          **kwargs)
  
      return handler.barcol, handler.opencol, handler.closecol
  
  
  class LineOnClosePlotHandler(object):
      legend_closes = [0.00, 0.66, 0.33, 1.00]
  
      def __init__(self,
                   ax, x, closes, color='k',
                   width=1, alpha=1.0,
                   label='_nolegend',
                   **kwargs):
  
          self.color = color
          self.alpha = alpha
  
          self.loc, = self.barcollection(
              x, closes,
              width=width,
              label=label,
              **kwargs)
  
          # add collections to the axis and return them
          ax.add_line(self.loc)
  
          # Update the axis
          ax.update_datalim(((x[0], min(closes)), (x[-1], max(closes))))
          ax.autoscale_view()
  
          # Add self as legend handler for this object
          mlegend.Legend.update_default_handler_map({self.loc: self})
  
      def legend_artist(self, legend, orig_handle, fontsize, handlebox):
          x0 = handlebox.xdescent
          y0 = handlebox.ydescent
          width = handlebox.width / len(self.legend_closes)
          height = handlebox.height
  
          # Generate the x axis coordinates (handlebox based)
          xs = [x0 + width * (i + 0.5) for i in range(len(self.legend_closes))]
  
          linecol, = self.barcollection(
              xs, self.legend_closes,
              width=1.5,
              scaling=height, bot=y0)
  
          linecol.set_transform(handlebox.get_transform())
          handlebox.add_artist(linecol)
  
          return linecol,
  
      def barcollection(self,
                        xs, closes,
                        width,
                        label='_nolegend',
                        scaling=1.0, bot=0,
                        **kwargs):
  
          # Prepack different zips of the series values
          scaled = [close * scaling + bot for close in closes]
  
          loc = mlines.Line2D(
              xs, scaled,
              color=self.color,
              lw=width,
              label=label,
              alpha=self.alpha,
              **kwargs)
  
          return loc,
  
  
  def plot_lineonclose(ax, x, closes,
                       color='k',
                       width=1.5,
                       alpha=1.0,
                       label='_nolegend',
                       **kwargs):
  
      handler = LineOnClosePlotHandler(
          ax, x, closes,
          color=color, width=width,
          alpha=alpha, label=label,
          **kwargs)
  
      return handler.loc,
  