import uuid
import pandas as pd
import maya
from jamboree import Jamboree, DBHandler

from loguru import logger
from typing import List, Dict, Any
from cytoolz.itertoolz import pluck
from crayons import yellow, blue, green, magenta
from link_jam_handle.data.pricing.simulated import SimulatedPriceHandler
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()