import uuid
import pandas as pd
import maya
from jamboree import Jamboree, DBHandler
# from contextlib import ContextDecorator
from typing import List, Dict, Any
from crayons import yellow, blue, green, magenta
from loguru import logger
from link_jam_handle.data.pricing.simulated import SimulatedPriceHandler
from cytoolz.itertoolz import pluck
from mlfinlab.portfolio_optimization.mean_variance import MeanVarianceOptimisation

"""
Note: Will remove soon. Going through a concept.
"""




class AllocationHandler(DBHandler):
    """Abstract handler that we use to keep track of information.
    """

    def __init__(self, limit=500, spacing=5, unit="hours"):
        super().__init__()
        self.entity = "allocation"
        self.required = {
            "user_id": str,
            "exchange": str,
            "episode": str,
            "live": bool
        }
        self._balance = 0
        self._limit = limit
        self._spacing = spacing
        self._unit = unit
        self.mean_variance_opt = MeanVarianceOptimisation()
    
    
        
    
    # -------------------------------------------------------------------
    # --------------------- Properties & Setters ------------------------
    # -------------------------------------------------------------------
    


    @property
    def limit(self):
        return self._limit
    
    @limit.setter
    def limit(self, limit):
        self._limit = limit
    

    @property
    def spacing(self):
        return self._spacing

    @spacing.setter
    def spacing(self, space):
        self._spacing = space



    @property
    def unit(self):
        return self._unit
    
    @unit.setter
    def unit(self, _unit):
        self._unit = _unit


    @property
    def assets(self) -> list:
        """ Get all of the assets that we're watching for the exchange """
        return self.latest_asset().get("assets", [])

    @property
    def allocation(self):
        return self.latest_allocation().get("allocation", {})
    
    @property
    def current_step(self):
        alt = {"detail": "steps"}
        count = self.count(alt=alt)
        return count
    
    @property
    def last_time(self):
        _late = self.latest_timestamp()
        
        return self.latest_timestamp().get("epoch", maya.now()._epoch)

    @property
    def should_rebalance(self) -> bool:
        """ Depending on the scenario, we determine if we should rebalance."""
        """
            if live:
                check if the determined amount of time has passed since we've rebalanced
            else:
                check if the predetermined number of steps have happened.
        """
        if self["live"]:
            current_epoch = maya.now()._epoch
            expected = maya.MayaDT(self.last_time).add(**{f"{self.unit}":self.spacing})._epoch
            # print(expected, current_epoch)
            if current_epoch >= expected:
                
                return True
        else:
            if self.current_step % self.spacing == 0:
                logger.info(green("Yes, we're going to rebalance", bold=True))
                return True
        return False

    def latest_asset(self):
        alt = {"detail": "assets"}
        asset = self.last(alt=alt)
        return asset
    
    def latest_allocation(self):
        alt = {"detail": "allocation"}
        asset = self.last(alt=alt)
        return asset
    

    def latest_timestamp(self):
        alt = {"detail": "timestamp"}
        asset = self.last(alt=alt)
        return asset


    def timestamp_count(self) -> int:
        """ Get the count of the asset adds we have. Not the number of assets. """
        alt = {"detail": "timestamp"}
        count = self.count(alt=alt)
        return count

    
    def asset_count(self) -> int:
        """ Get the count of the asset adds we have. Not the number of assets. """
        alt = {"detail": "assets"}
        count = self.count(alt=alt)
        return count
    


    def count_steps(self):
        alt = {"detail": "steps"}
        count = self.count(alt=alt)
        return count


    def add_assets(self, assets:List[str]):
        """ Add a list of assets to monitor. Used for tracking multiple assets inside of a portfolio. """
        if len(assets) == 0:
            return
        asset_list = self.assets
        for asset in assets:
            if asset not in asset_list:
                asset_list.append(asset)
        self.save_monitored_assets(asset_list)

    def add_asset(self, asset_name:str):
        """ Add an asset that will be monitored in an exchange"""
        if asset_name == "":
            return
        asset_list = self.assets
        asset_list.append(asset_name)
        self.save_monitored_assets(asset_list)
        
    
    def save_monitored_assets(self, assets:list):
        """ Save monitored assets """
        alt = {"detail": "assets"}
        monitored = {
            "assets": assets
        }
        self.save(monitored, alt=alt)
    
    def save_allocation(self, allocation:dict):
        """ Save monitored assets """
        alt = {"detail": "allocation"}
        monitored = {
            "allocation": allocation
        }
        self.save(monitored, alt=alt)
    
    def save_timestamp(self, _time=None):
        """ Save monitored assets """
        if _time is None:
            _time = maya.now()._epoch
        alt = {"detail": "timestamp"}
        monitored = {
            "epoch": _time
        }
        self.save(monitored, alt=alt)


    def save_step(self):
        alt = {"detail": "steps"}
        monitored = {}
        self.save(monitored, alt=alt)
    







    def _filter_assets(self, all_assets:dict) -> dict:
        assets = self.assets
        users_portfolio_to_rebalance = {}
        _all_assets = all_assets.get("assets", None)
        if _all_assets is None:
            return {}
        for asset in assets:
            current_asset = _all_assets.get(asset, None)
            if current_asset is not None:
                users_portfolio_to_rebalance[asset] = current_asset

        return users_portfolio_to_rebalance
    

    def _extract_frame(self, portfolio_assets:dict):
        time_is_filled = False
        time_index = []
        portfolio_close = {}
        for asset in portfolio_assets.keys():
            current_pricing = portfolio_assets.get(asset, [])
            portfolio_close[asset] = list(pluck("close", current_pricing))
            if time_is_filled == False:
                time_index = list(pluck("time", current_pricing))
                time_is_filled = True
        
        pandas_index = pd.to_datetime(time_index, unit='s')
        pclose = pd.DataFrame(portfolio_close, index=pandas_index)
        return pclose

    

    


    def _finish_step(self, is_rebalance=False):
        """ Finish update that we're not """
        if is_rebalance == True:
            # Give a timestamp of the last rebalance.
            print(yellow("saving timestamp"))
            self.save_timestamp()
        self.save_step()
        


    def step(self, all_assets:Dict[str, Any]) -> None:
        """ Accept all of the assets for a given allocation and begin"""
        

        if self.should_rebalance == True:
            users_portfolio_to_rebalance = self._filter_assets(all_assets)

            portfolio_keys = users_portfolio_to_rebalance.keys()
            if len(portfolio_keys) == 0:
                logger.info(magenta("There are no portfolio keys", bold=True))
                self._finish_step()
                return

            pclose = self._extract_frame(users_portfolio_to_rebalance)
            self.mean_variance_opt.allocate(pclose)
            current_allocation = self.mean_variance_opt.weights.to_dict('records')[0]
            logger.info(current_allocation)
            self.save_allocation(current_allocation)
            self._finish_step(is_rebalance=True)
            return
            # Set timestamp aka, the last time we did a rebalance.
        self._finish_step()


    def _reset_assets(self) -> None:
        if self.asset_count() == 0:
            logger.info(yellow("Setting the first assets"))
            asset_list = self.assets
            self.save_monitored_assets(asset_list)
        else:
            logger.info(blue("Already reset"))

    def _reset_timestamp(self):
        if self.timestamp_count() == 0:
            logger.info("We're setting a null timestamp. This way the rebalance will happen automatically upon the first step.")
            self.save_timestamp(_time=maya.when('2 years ago')._epoch)


    def reset(self):
        self._reset_assets()
        self._reset_timestamp()


def experiment_main():
    jam = Jamboree()
    allo_handler = AllocationHandler(unit="seconds", spacing=2)
    price_handler = SimulatedPriceHandler()

    allo_handler.event = jam
    price_handler.event = jam
    
    
    
    allo_handler["user_id"] = "1"
    allo_handler["episode"] = uuid.uuid4().hex
    allo_handler["exchange"] = f"binance-{uuid.uuid4().hex}"
    allo_handler["live"] = True
    
    asset_list = ["BTC", "ATL", "TRX", "ETH", "BCH", "XRP", "LTC", "EOS", "ADA", "XMR", "LINK", "HT"]
    # This live variable locks certain commands, and allows us to query faster. 
    price_handler['episode'] = uuid.uuid4().hex # set this to live if we're trying to use live prices.
    price_handler['exchange'] = "binance"
    price_handler['live'] = False
    price_handler['obs_num'] = 2000
    price_handler.add_assets(asset_list)
    price_handler.reset()


    for _ in range(10):
        allo_handler.reset()
    allo_handler.add_assets(["BTC", "ETH", "TRX", "BCH"])


    while True:

        obs = price_handler.step()

        # Get the observation number
        i = obs.get("step_num", 0)
        assets = obs.get("assets", {})
        done = obs.get("done", False)
        if done:
            logger.info(blue("Observation Complete"))
            break
        # print()
        # latest_price = price_handler.latest_price("BTC")
        # logger.info(yellow(latest_price['time']))
        # logger.info(red(price_handler.current_time))

        allo_handler.step(obs)
    

if __name__ == "__main__":
    experiment_main()