import math
from stock_info import StockSceneInfo
from stock_info import StockHoldInfo
import const as cst
import numpy as np


class StockHolder():

    def __init__(self,is_safe,verbose = False,single_lock=True,max_stock_day=20) -> None:
        self.money = cst.INITIAL_CAPITAL
        self.stock_info = dict()
        self.is_risk = False
        self.safe_money = 0
        self.is_safe = is_safe
        self.single_lock = single_lock
        self.max_stock_day = max_stock_day
        self.verbose = verbose
        self.once_money = cst.ONCE_DEAL_CAPIAL

    def update(self,code,last_info) :
        if code in self.stock_info:
            self.stock_info[code].SceneInfo = last_info
            self.stock_info[code].HoldDay+=1
            if self.stock_info[code].HoldDay > self.max_stock_day :
                price = last_info.Open # 核按钮，开盘就卖出
                self.sell(code,price)    

    def buy(self,code,stock:StockSceneInfo,price:np.float64):

        if price<stock.Low :
            return False
        if price>stock.Open:
            price = stock.Open
        buy_stock_count = math.floor( self.once_money / (price * 100))
        need_money = buy_stock_count*100*price
        if self.money < need_money:
            return False
        if self.single_lock and code in self.stock_info:
            return False
        if code in self.stock_info:
            self.stock_info[code].Count += buy_stock_count
        else:
            self.stock_info[code] = StockHoldInfo(code,stock,buy_stock_count,0)  

        self.money = self.money - need_money
        
        if self.verbose:
            print(f'{stock.Date} BUY : {code} {price} {self.get_fortune()}')
        return True

    
    def sell(self,code,price:np.float64):

        if code not in self.stock_info :
            return False
        sell_info = self.stock_info[code]
        if price > sell_info.SceneInfo.High :
            return False
        if price < sell_info.SceneInfo.Open:
            price = sell_info.SceneInfo.Open
        gain_money = sell_info.Count * price * 100
        self.money += gain_money*0.998 #扣除手续费
        
        del self.stock_info[code]
        if self.is_safe and self.money>cst.INITIAL_CAPITAL:
            self.safe_money += self.money-cst.INITIAL_CAPITAL
            self.money = cst.INITIAL_CAPITAL
        if self.verbose:
            print(f'{sell_info.SceneInfo.Date} SELL : {code} {price} {self.get_fortune()}')

        return True


    def get_fortune(self):
        money = self.money + self.safe_money
        for key,value in self.stock_info.items () :
            stock:StockHoldInfo = value
            money+=(stock.Count * stock.SceneInfo.Standard * 100)
        return money