from typing import Dict

import numpy as np
import pandas as pd


class Stock:

    def __init__(self, symbol, sources: Dict[str, pd.DataFrame]):
        self.symbol = symbol
        self.sources = sources
        self.generate()

    def generate(self):
        for interval in self.sources.keys():
            self._generate_interval(interval)

    def _generate_interval(self, interval):
        df = self.sources[interval]
        if df is not None and not df.empty:
            df.columns = df.columns.str.lower()
            df.index.name = 'datetime'
            df = self._calculate_ma(df)
            df = self._calculate_macd(df)
            df = self._calculate_bb(df)
        self.sources[interval] = df

    @staticmethod
    def _calculate_ma(df):
        for window in (5, 10, 20, 30, 60, 120, 250):
            df[f'ma{window}'] = df['close'].rolling(window=window).mean()
        return df

    @staticmethod
    def _calculate_macd(df):
        ema_fast = df['close'].ewm(span=12, adjust=False).mean()
        ema_slow = df['close'].ewm(span=26, adjust=False).mean()
        df['macd'] = ema_fast - ema_slow
        df['macd_dea'] = df['macd'].ewm(span=9, adjust=False).mean()
        df['macd_dif'] = df['macd'] - df['macd_dea']
        return df

    @staticmethod
    def _calculate_bb(df):
        df['bb'] = df['close'].rolling(window=20).mean()
        rolling_std = df['close'].rolling(window=20).std()
        df['bb_upper'] = df['bb'] + (2 * rolling_std)
        df['bb_lower'] = df['bb'] - (2 * rolling_std)
        return df

    @staticmethod
    def is_top(item):
        return np.isnan(item.low) and not np.isnan(item.high)

    @staticmethod
    def is_bottom(item):
        return np.isnan(item.high) and not np.isnan(item.low)

    @staticmethod
    def to_dataframe(rows, columns=None):
        if rows is None or len(rows) == 0:
            return None
        df = pd.DataFrame(rows).set_index('Index')
        df.index.name = 'datetime'
        return df if columns is None else df[columns]
