from stockdatadownloader.CommonVariable import STOCK_BIT_CSV_FILE_PATH
from stockdatadownloader.StockBitColumn import MONTH_COLUMN, DATE_COLUMN
from stockdatadownloader.fetcher.StockBitFetcher import StockBitFetcher
import pandas as pd


class StockBitQuery:
    list_df = {}

    @staticmethod
    def get_trading_dates(stock_name):
        df = pd.read_csv(STOCK_BIT_CSV_FILE_PATH + stock_name + '_price.csv')
        return df[DATE_COLUMN]

    # all date are string in format yyyy-mm-dd
    @staticmethod
    def query_data_fundamental(stock_name, date, column_name):
        quarter = StockBitQuery.__get_quarter_string(date)

        for report_type in StockBitFetcher.REPORT_TYPES:
            file_name = STOCK_BIT_CSV_FILE_PATH + stock_name + '_' + report_type + '.csv'
            df = StockBitQuery.__get_csv_file(file_name)

            if df is None:
                return

            if column_name in df:
                response = df.query(MONTH_COLUMN + " == '" + quarter + "'")[column_name].values
                if len(response) == 1:
                    return response[0]

        return None

    # all date are string in format yyyy-mm-dd
    @staticmethod
    def query_data_price(stock_name, date, column_name):
        file_name = STOCK_BIT_CSV_FILE_PATH + stock_name + '_price.csv'
        df = StockBitQuery.__get_csv_file(file_name)

        if df is None:
            return

        if column_name in df:
            response = df.query(DATE_COLUMN + " == '" + date + "'")[column_name].values
            if len(response) == 1:
                return response[0]

        return None

    @staticmethod
    def __get_csv_file(file_name):
        if file_name in StockBitQuery.list_df:
            return StockBitQuery.list_df[file_name]

        try:
            df = pd.read_csv(file_name)
        except FileNotFoundError:
            df = None

        StockBitQuery.list_df[file_name] = df
        return df

    @staticmethod
    def __get_quarter_string(date_string):
        date = pd.to_datetime(date_string)

        quarter = StockBitQuery.__get_quarter(date.month)
        year = date.year

        if quarter == 0:
            quarter = 4
            year -= 1

        response = "Q" + str(quarter) + " " + str(year)
        return response

    @staticmethod
    def __get_quarter(month):
        return (month - 1) // 3
