from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
from datetime import datetime
import pandas as pd


def HA(o=[], h=[], l=[], c=[], v=[], ts=[]):
    ha = None
    try:
        inputs = {
            'open': o,
            'high': h,
            'low': l,
            'close': c,
            'volume': v,
            'ts': ts
        }

        df = pd.DataFrame(
            data=inputs,
            columns=['open', 'high', 'low', 'close', 'volume', 'ts'])
        df = df.set_index('ts')

        ha = _heikin_ashi(df)
    except Exception as e:
        print(datetime.now().strftime("%Y%m%d %H:%M:%S"), "ERROR",
              "_heikin_ashi ", 'Exception', str(e))

    return ha


def _heikin_ashi(df):
    '''
    @brief  heikin ashi
            参考 https://github.com/songv1/2019-Cryptocurrency-Trading-Algorithm/blob/master/crypto_heikinashi_charts.py
    '''
    heikin_ashi_df = pd.DataFrame(index=df.index.values,
                                  columns=['open', 'high', 'low', 'close'])
    heikin_ashi_df['close'] = (df['open'] + df['high'] + df['low'] +
                               df['close']) / 4
    for i in range(len(df)):
        if i == 0:
            heikin_ashi_df.iat[0, 0] = df['open'].iloc[0]
        else:
            heikin_ashi_df.iat[i, 0] = (heikin_ashi_df.iat[i - 1, 0] +
                                        heikin_ashi_df.iat[i - 1, 3]) / 2
    heikin_ashi_df['high'] = heikin_ashi_df.loc[:, ['open', 'close']].join(
        df['high']).max(axis=1)
    heikin_ashi_df['low'] = heikin_ashi_df.loc[:, ['open', 'close']].join(
        df['low']).min(axis=1)
    return heikin_ashi_df
