import torch
import numpy as np
import torch.nn.functional as F
import math


class StockAlgo:
    @staticmethod
    def loss(output, pct_change, short_filter: None | str = None):
        raise NotImplementedError

    @staticmethod
    def metric(output, pct_change, short_filter: None | str = None):
        raise NotImplementedError

    @staticmethod
    def accumulate(output, pct_change, short_filter: None | str = None):
        raise NotImplementedError


def get_stock_algo(target_type, stock_loss_mode) -> StockAlgo:
    if target_type != "logpctchange" and target_type != "pctchange":
        raise Exception(f"Invalid Target Type: {target_type}")

    stock_algo = stock_loss_mode.split("-")[0]

    if "tanh" in stock_algo:
        if stock_algo == "tanh":
            assert target_type == "pctchange"
            return PctProfitTanh

        assert target_type == "logpctchange"

        if stock_algo == "tanhv1":
            return LogPctProfitTanhV1
        elif stock_algo == "tanhv2":
            return LogPctProfitTanhV2
        elif stock_algo == "tanhv3":
            return LogPctProfitTanhV3
        elif stock_algo == "tanhv4":
            return LogPctProfitTanhV4

        raise Exception("Invalid tanh loss")
    elif "dir" == stock_algo:
        return (
            LogPctProfitDirection
            if target_type == "logpctchange"
            else PctProfitDirection
        )

    raise Exception(f"Invalid Stock Loss Mode: {stock_loss_mode}")


def get_short_filter(stock_loss_mode: str) -> None | str:
    stock_algo_split = stock_loss_mode.split("-")
    if len(stock_algo_split) == 1:
        return None
    else:
        assert stock_algo_split[-1] in ["ns", "os"]
        return stock_algo_split[-1]


def apply_short_filter(output, raw, short_filter: None | str):
    if short_filter is None:
        return raw
    elif short_filter == "ns":
        return raw[output > 0]
    elif short_filter == "os":
        return raw[output < 0]
    raise Exception(f"Invalid short filter: {short_filter}")


def apply_threshold_metric(output, other, threshold=0.0002):
    output_tresh = output[np.abs(output) >= threshold]
    other = other[np.abs(output) >= threshold]
    return output_tresh, other


def pct_direction(output, pct_change):
    return np.sum(np.sign(output) == np.sign(pct_change)) / len(pct_change)


def pct_direction_torch(output, pct_change):
    output = torch.tanh(output / 2)
    return torch.sum(torch.sign(output) == torch.sign(pct_change)) / len(pct_change)
    #  y = torch.tanh(output / 2)
    #  raw = torch.log(1 + y * (torch.exp(log_pct_change) - 1))
    #  return torch.exp(raw.sum())


class PctProfitDirection(StockAlgo):
    """
    Percent profit with investing everything strategy.
    """

    @staticmethod
    def loss(output, pct_change, short_filter: None | str = None):
        raw = pct_change * torch.sign(output) + 1
        return apply_short_filter(output, raw, short_filter)

    @staticmethod
    def metric(output, pct_change, short_filter: None | str = None):
        raw = pct_change * np.sign(output) + 1
        return apply_short_filter(output, raw, short_filter).prod()

    @staticmethod
    def accumulate(output, pct_change, short_filter: None | str = None):
        raw = pct_change * np.sign(output) + 1
        return np.cumprod(apply_short_filter(output, raw, short_filter))


class LogPctProfitDirection(StockAlgo):
    """
    Percent profit with investing everything strategy.
    """

    @staticmethod
    def loss(output, log_pct_change, short_filter: None | str = None):
        raw = log_pct_change * torch.sign(output)
        return apply_short_filter(output, raw, short_filter)

    @staticmethod
    def metric(output, log_pct_change, short_filter: None | str = None):
        raw = log_pct_change * np.sign(output)
        return np.exp(apply_short_filter(output, raw, short_filter).sum())

    @staticmethod
    def accumulate(output, log_pct_change, short_filter: None | str = None):
        raw = log_pct_change * np.sum(output)
        return np.exp(np.cumsum(apply_short_filter(output, raw, short_filter)))


class PctProfitTanh(StockAlgo):
    """
    Percent profit with investing tanh partial purchase
    """

    @staticmethod
    def loss(output, pct_change, short_filter: None | str = None):
        raw = (pct_change * torch.tanh(output)) + 1
        return apply_short_filter(output, raw, short_filter)

    @staticmethod
    def metric(output, pct_change, short_filter: None | str = None):
        raw = pct_change * np.tanh(output) + 1
        return apply_short_filter(output, raw, short_filter).prod()

    @staticmethod
    def accumulate(output, pct_change, short_filter: None | str = None):
        raw = pct_change * np.tanh(output) + 1
        return np.cumprod(apply_short_filter(output, raw, short_filter))


class LogPctProfitTanhV1(StockAlgo):
    """
    Percent profit with investing tanh partial purchase

    V1: just uses a tanh based multiplier. This is inaccurate as the bounds for shorting is not -1 but `log(1-pctchange)/log(1+pctchange)`.
    However this quantity is near -1.
    """
    @staticmethod
    #  def loss(output, log_pct_change, short_filter: None | str = None):
        #  assert not torch.isnan(output).any(), "output is nan"
        #  tanh = torch.tanh(output)
        #  raw = log_pct_change * tanh
        #  return apply_short_filter(output, raw, short_filter)
    #  def loss(output, log_pct_change, short_filter: None | str = None):
        #  assert not torch.isnan(output).any(), "output is nan"
        #  same_sign = (output * log_pct_change) >= 0
        #  # 对的样本：-tanh(o)*y
        #  loss_good = torch.tanh(output) * log_pct_change
        #  # 错的样本：线性惩罚 -o * y
        #  loss_bad = output * log_pct_change
        #  # 合并
        #  raw = torch.where(same_sign, loss_good, loss_bad)
        #  return apply_short_filter(output, raw, short_filter)
    #  def loss(output, log_pct_change, short_filter: None | str = None):
        #  assert not torch.isnan(output).any(), "output is nan"
        #  same_sign = log_pct_change >= 0
        #  # 稳定计算 log-sigmoid
        #  log_pos = -F.softplus(-output)  # log(sigmoid(x))
        #  log_neg = -output - F.softplus(-output)  # log(1 - sigmoid(x))
        #  #  gamma = 2
        #  #  log_pos = (1-torch.sigmoid(output)).pow(gamma) * log_pos
        #  #  log_neg = torch.sigmoid(output).pow(gamma) * log_neg
        #  raw = torch.where(same_sign, log_pos, log_neg) * torch.abs(log_pct_change)
        #  raw = torch.sum(raw) / torch.sum(torch.abs(log_pct_change))
        #  return apply_short_filter(output, raw, short_filter)
    def loss(output, log_pct_change, short_filter: None | str = None):
        assert not torch.isnan(output).any(), "output is nan"
        #  print(log_pct_change.mean(), log_pct_change.std())
        #  while 1:pass
        x = 3000 * log_pct_change
        y = torch.sigmoid(x)
        #  y = torch.special.expit(x)
        #  print(y)
        log_pos = -F.softplus(-output)  # log(sigmoid(x))
        log_neg = -output - F.softplus(-output)  # log(1 - sigmoid(x))
        gt_loss = -(y*F.softplus(-x) + (1-y)*(x+F.softplus(-x)))
        #  gt_loss = y*torch.log(y) + (1-y)*torch.log(1-y)
        entropy_weight = 1 + gt_loss / torch.log(torch.tensor(2.0))
        raw = y * log_pos + (1-y) * log_neg - gt_loss
        raw = entropy_weight * raw
        return apply_short_filter(output, raw, short_filter)

    @staticmethod
    def sharpe(output, log_pct_change):
        p = torch.tanh(output/2)
        position, _ = LogPctProfitTanhV1.postion_rule(p, L=1e1, clamp=False)
        pnl =  position * (torch.exp(log_pct_change) - 1)
        mu  = pnl.mean()
        sigma = pnl.std(unbiased=False) + 1e-18
        sharpe = mu / sigma
        return sharpe

    @staticmethod
    def postion_rule(p, input_scale=None, L=1e0, clamp=True):
        if input_scale is None:
            scale = p.abs().mean() + 1e-18
        else:
            scale = input_scale
        if not clamp:
            return (p / scale) * L, scale
        else:
            return (p / scale) * L, scale
            #  return torch.clamp(p / scale, min=-1.0, max=1.0) * L, scale
        #  if rule == 0:
            #  tau = 0.01           # 阈值，可以调
            #  w = torch.zeros_like(p)
            #  w[p >  tau] =  L    # 做多
            #  w[p < -tau] = -L   # 做空
            #  return w
        #  elif rule == 1:
            #  scale = p.abs().mean() + 1e-18     # 平均强度做归一化
            #  return (p / scale) * L               # 最终仓位
        #  elif rule == 2:
            #  return p
        #  else:
            #  return p * L

    @staticmethod
    def metric(output, log_pct_change, short_filter: None | str = None, input_scale=None):
        #  print(log_pct_change.shape, output.shape, log_pct_change.std(), log_pct_change.mean())
        #  while 1:pass
        # TODO: Potentially clip the negative tanh outputs: max(tanh, log(1-pctchange)/log(1+pctchange))
        if isinstance(output, np.ndarray):
            output = torch.tensor(output)
            log_pct_change = torch.tensor(log_pct_change)
        y = torch.tanh(output / 2)
        postion, scale = LogPctProfitTanhV1.postion_rule(y, input_scale)
        raw = torch.log(1 + postion * (torch.exp(log_pct_change) - 1))
        return apply_short_filter(output, raw, short_filter), scale

    @staticmethod
    def accumulate(output, log_pct_change, short_filter: None | str = None, input_scale=None):
        # TODO: Potentially clip the negative tanh outputs: max(tanh, log(1-pctchange)/log(1+pctchange))
        if isinstance(output, np.ndarray):
            output = torch.tensor(output)
            log_pct_change = torch.tensor(log_pct_change)
        y = torch.tanh(output / 2)
        postion, scale = LogPctProfitTanhV1.postion_rule(y, input_scale)
        raw = torch.log(1 + postion * (torch.exp(log_pct_change) - 1))
        raw = raw.detach().cpu().numpy()
        if not isinstance(scale, float):
            scale = scale.item()
        return np.exp(np.cumsum(apply_short_filter(output, raw, short_filter))), scale

    #  @staticmethod
    #  def metric(output, log_pct_change, short_filter: None | str = None):
        #  y = torch.tanh(output)
        #  raw = torch.log(1 + y * (torch.exp(log_pct_change) - 1))
        #  #  raw = log_pct_change * y
        #  return torch.exp(torch.sum(apply_short_filter(output, raw, short_filter)))

    #  @staticmethod
    #  def metric(output, log_pct_change, short_filter: None | str = None):
        #  # TODO: Potentially clip the negative tanh outputs: max(tanh, log(1-pctchange)/log(1+pctchange))
        #  tanh = np.tanh(output)
        #  raw = log_pct_change * tanh
        #  return np.exp(apply_short_filter(output, raw, short_filter).sum())

    #  @staticmethod
    #  def accumulate(output, log_pct_change, short_filter: None | str = None):
        #  # TODO: Potentially clip the negative tanh outputs: max(tanh, log(1-pctchange)/log(1+pctchange))
        #  tanh = np.tanh(output)
        #  raw = log_pct_change * tanh
        #  return np.exp(np.cumsum(apply_short_filter(output, raw, short_filter)))


class LogPctProfitTanhV2(StockAlgo):
    """
    Percent profit with investing tanh partial purchase

    V2: Scales the just negative tanh output so that they are between `log(1-pctchange)/log(1+pctchange)` and  `0`.
    """

    @staticmethod
    def loss(output, log_pct_change, short_filter: None | str = None):
        # The partial purchase multiplier is from [log(1-pctchange)/log(1+pctchange), 1]
        # mult_min is multi_min log(1-pctchange)/log(1+pctchange) and is negative
        # TODO: Look into logaddexp
        pct_change_mult = torch.exp(log_pct_change)
        mult_min = torch.log(-pct_change_mult + 2) / log_pct_change
        mult_min[mult_min != mult_min] = -1  # get rid of nan from 0/0

        tanh = torch.tanh(output)

        # Scale only the negative side of the tanh
        mult_min[tanh >= 0] = -1.0
        scaled_tanh = tanh * (-mult_min)

        raw = log_pct_change * scaled_tanh
        return apply_short_filter(output, raw, short_filter)

    @staticmethod
    def metric(output, log_pct_change, short_filter: None | str = None):
        pct_change_mult = np.exp(log_pct_change)
        mult_min = np.log(-pct_change_mult + 2) / log_pct_change
        mult_min[mult_min != mult_min] = -1

        tanh = np.tanh(output)

        mult_min[tanh >= 0] = -1.0
        scaled_tanh = tanh * (-mult_min)

        raw = log_pct_change * scaled_tanh
        return np.exp(apply_short_filter(output, raw, short_filter).sum())

    @staticmethod
    def accumulate(output, log_pct_change, short_filter: None | str = None):
        pct_change_mult = np.exp(log_pct_change)
        mult_min = np.log(-pct_change_mult + 2) / log_pct_change
        mult_min[mult_min != mult_min] = -1

        tanh = np.tanh(output)

        mult_min[tanh >= 0] = -1.0
        scaled_tanh = tanh * (-mult_min)

        raw = log_pct_change * scaled_tanh
        return np.exp(np.cumsum(apply_short_filter(output, raw, short_filter)))


class LogPctProfitTanhV3(StockAlgo):
    """
    Percent profit with investing tanh partial purchase

    V3: Scales the whole tanh output so that they are between `log(1-pctchange)/log(1+pctchange)` and  `1`.
    Note: We are most likely going to remove this version because due to the shift `output>0` will not necessarily mean buy, same with `output<0` for short
    """

    @staticmethod
    def loss(output, log_pct_change, short_filter: None | str = None):
        # The partial purchase multiplier is from [log(1-pctchange)/log(1+pctchange), 1]
        # mult_min is multi_min log(1-pctchange)/log(1+pctchange) and is negative
        # Look into logaddexp
        pct_change_mult = torch.exp(log_pct_change)
        mult_min = torch.log(-pct_change_mult + 2) / log_pct_change
        mult_min[mult_min != mult_min] = -1  # get rid of nan from 0/0

        tanh = torch.tanh(output)

        # Could just use a sigmoid with 2*output
        at_sigmoid_bounds = (tanh + 1) / 2.0
        raw = log_pct_change * ((1 - mult_min) * at_sigmoid_bounds + mult_min)

        return apply_short_filter(output, raw, short_filter)

    @staticmethod
    def metric(output, log_pct_change, short_filter: None | str = None):
        pct_change_mult = np.exp(log_pct_change)
        mult_min = np.log(-pct_change_mult + 2) / log_pct_change
        mult_min[mult_min != mult_min] = -1

        tanh = np.tanh(output)

        at_sigmoid_bounds = (tanh + 1) / 2.0
        raw = log_pct_change * ((1 - mult_min) * at_sigmoid_bounds + mult_min)

        return np.exp(apply_short_filter(output, raw, short_filter).sum())

    @staticmethod
    def accumulate(output, log_pct_change, short_filter: None | str = None):
        pct_change_mult = np.exp(log_pct_change)
        mult_min = np.log(-pct_change_mult + 2) / log_pct_change
        mult_min[mult_min != mult_min] = -1

        tanh = np.tanh(output)

        at_sigmoid_bounds = (tanh + 1) / 2.0
        raw = log_pct_change * ((1 - mult_min) * at_sigmoid_bounds + mult_min)

        return np.exp(np.cumsum(apply_short_filter(output, raw, short_filter)))


class LogPctProfitTanhV4(StockAlgo):
    """
    Percent profit with investing tanh partial purchase and 2x multiplier on losses

    V4: just uses a tanh based multiplier. This is inaccurate as a metric but fine as a loss.
    """

    @staticmethod
    def loss(output, log_pct_change, short_filter: None | str = None):
        assert not torch.isnan(output).any(), "output is nan"
        tanh = torch.tanh(output)
        raw = log_pct_change * tanh

        # The 1e-8 is because torch.sqrt(0) doesn't have a gradient
        raw_sqrt = 2 * raw  # torch.sqrt(torch.abs(raw) + 1e-8)
        is_loss = raw < 0
        is_gain = raw >= 0

        raw_w_direction_punishment = raw * is_gain + raw_sqrt * is_loss

        return apply_short_filter(output, raw_w_direction_punishment, short_filter)

    @staticmethod
    def metric(output, log_pct_change, short_filter: None | str = None):
        raise Exception("LogPctProfitTanhV4 is not a valid metric")
        # # TODO: Potentially clip the negative tanh outputs: max(tanh, log(1-pctchange)/log(1+pctchange))
        # tanh = np.tanh(output)
        # raw = log_pct_change * tanh
        # return np.exp(apply_short_filter(output, raw, short_filter).sum())

    @staticmethod
    def accumulate(output, log_pct_change, short_filter: None | str = None):
        raise Exception("LogPctProfitTanhV4 is not a valid metric")
        # # TODO: Potentially clip the negative tanh outputs: max(tanh, log(1-pctchange)/log(1+pctchange))
        # tanh = np.tanh(output)
        # raw = log_pct_change * tanh
        # return np.exp(np.cumsum(apply_short_filter(output, raw, short_filter)))