from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal
from math import ceil, floor, log2
from typing import Union

import torch
from ppq.core import RoundingPolicy


def ppq_numerical_round(value: float, 
    policy: RoundingPolicy=RoundingPolicy.ROUND_HALF_EVEN) -> int:
    """
        reference: https://en.wikipedia.org/wiki/Rounding
        
        decimal defination:
            - decimal.ROUND_CEILING (towards Infinity)
            - decimal.ROUND_DOWN (towards zero)
            - decimal.ROUND_FLOOR (towards -Infinity)
            - decimal.ROUND_HALF_DOWN (to nearest with ties going towards zero)
            - decimal.ROUND_HALF_EVEN (to nearest with ties going to nearest even integer)
            - decimal.ROUND_HALF_UP (to nearest with ties going away from zero)
            - decimal.ROUND_UP (away from zero)
            - decimal.ROUND_05UP (away from zero if last digit after rounding towards zero would have been 0 or 5; otherwise towards zero)

    Args:
        value (float): [description]
        policy (RoundingPolicy, optional): [description]. Defaults to RoundingPolicy.ROUND_HALF_EVEN.

    Raises:
        ValueError: [description]

    Returns:
        int: [description]
    """
    assert isinstance(value, float), 'numerical round only takes effect on float number.'
    if policy == RoundingPolicy.ROUND_HALF_EVEN:
        return int(Decimal(value).quantize(exp=Decimal(1), rounding=ROUND_HALF_EVEN))
    elif policy == RoundingPolicy.ROUND_HALF_UP:
        if value > 0: return int(Decimal(value).quantize(exp=Decimal(1), rounding=ROUND_HALF_UP))
        else: return int(Decimal(value).quantize(exp=Decimal(1), rounding=ROUND_HALF_DOWN))
    elif policy == RoundingPolicy.ROUND_HALF_DOWN:
        if value > 0: return int(Decimal(value).quantize(exp=Decimal(1), rounding=ROUND_HALF_DOWN))
        else: return int(Decimal(value).quantize(exp=Decimal(1), rounding=ROUND_HALF_UP))
    elif policy == RoundingPolicy.ROUND_HALF_TOWARDS_ZERO:
        return ppq_numerical_round(value, RoundingPolicy.ROUND_HALF_DOWN)
    elif policy == RoundingPolicy.ROUND_HALF_FAR_FORM_ZERO:
        return ppq_numerical_round(value, RoundingPolicy.ROUND_HALF_UP)
    elif policy == RoundingPolicy.ROUND_TO_NEAR_INT:
        if value > 0: return floor(value + 0.5)
        else: return ceil(value - 0.5)
    elif policy == RoundingPolicy.ROUND_UP:
        return ceil(value)
    else:
        raise ValueError('Unexpected rounding policy found.')

def ppq_tensor_round(value: torch.Tensor, 
    policy:RoundingPolicy=RoundingPolicy.ROUND_HALF_EVEN) -> torch.Tensor:
    """
        reference: https://en.wikipedia.org/wiki/Rounding

    Args:
        value (torch.Tensor): [description]
        policy (RoundingPolicy, optional): [description]. Defaults to RoundingPolicy.ROUND_HALF_EVEN.

    Raises:
        ValueError: [description]

    Returns:
        torch.Tensor: [description]
    """
    assert isinstance(value, torch.Tensor), 'tensor round only takes effect on torch tensor.'
    if policy == RoundingPolicy.ROUND_HALF_EVEN:
        # default rounding policy of torch is ROUND_TO_NEAR_EVEN
        # try this: print(torch.Tensor([1.5, 2.5, 3.5, 4.5]).round())
        # However it may generate unexpected results due to version difference.
        return value.round()
    elif policy == RoundingPolicy.ROUND_UP:
        return value.ceil()
    elif policy == RoundingPolicy.ROUND_HALF_TOWARDS_ZERO:
        return torch.sign(value) * torch.ceil(value.abs() - 0.5)
    elif policy == RoundingPolicy.ROUND_HALF_FAR_FORM_ZERO:
        return torch.sign(value) * torch.floor(value.abs() + 0.5)
    elif policy == RoundingPolicy.ROUND_HALF_DOWN:
        return torch.ceil(value - 0.5)
    elif policy == RoundingPolicy.ROUND_HALF_UP:
        return torch.floor(value + 0.5)
    elif policy == RoundingPolicy.ROUND_TO_NEAR_INT:
        raise NotImplementedError(f'Torch Tensor can not use this rounding policy({policy}) try ROUND_HALF_EVEN instead.')
    else:
        raise ValueError('Unexpected rounding policy found.')

def ppq_round_to_power_of_2(value: Union[float, int], 
    policy: RoundingPolicy=RoundingPolicy.ROUND_UP) -> float:
    if value == 0: return 0
    sign = 1 if value >= 0 else -1
    assert isinstance(value, float) or isinstance(value, int), \
        'power-of-2 round only takes effect on float or int.'
    return sign * float(pow(2, ppq_numerical_round(log2(sign * value), policy=policy)))

[((36, 19, 36, 33), 'decimal.Decimal', 'Decimal', ({(36, 27, 36, 32): 'value'}, {}), '(value)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((36, 47, 36, 57), 'decimal.Decimal', 'Decimal', ({(36, 55, 36, 56): '(1)'}, {}), '(1)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((79, 15, 79, 32), 'torch.sign', 'torch.sign', ({(79, 26, 79, 31): 'value'}, {}), '(value)', False, 'import torch\n'), ((97, 51, 97, 69), 'math.log2', 'log2', ({(97, 56, 97, 68): '(sign * value)'}, {}), '(sign * value)', False, 'from math import ceil, floor, log2\n'), ((81, 15, 81, 32), 'torch.sign', 'torch.sign', ({(81, 26, 81, 31): 'value'}, {}), '(value)', False, 'import torch\n'), ((83, 15, 83, 38), 'torch.ceil', 'torch.ceil', ({(83, 26, 83, 37): '(value - 0.5)'}, {}), '(value - 0.5)', False, 'import torch\n'), ((38, 33, 38, 47), 'decimal.Decimal', 'Decimal', ({(38, 41, 38, 46): 'value'}, {}), '(value)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((38, 61, 38, 71), 'decimal.Decimal', 'Decimal', ({(38, 69, 38, 70): '(1)'}, {}), '(1)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((39, 25, 39, 39), 'decimal.Decimal', 'Decimal', ({(39, 33, 39, 38): 'value'}, {}), '(value)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((39, 53, 39, 63), 'decimal.Decimal', 'Decimal', ({(39, 61, 39, 62): '(1)'}, {}), '(1)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((85, 15, 85, 39), 'torch.floor', 'torch.floor', ({(85, 27, 85, 38): '(value + 0.5)'}, {}), '(value + 0.5)', False, 'import torch\n'), ((41, 33, 41, 47), 'decimal.Decimal', 'Decimal', ({(41, 41, 41, 46): 'value'}, {}), '(value)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((41, 61, 41, 71), 'decimal.Decimal', 'Decimal', ({(41, 69, 41, 70): '(1)'}, {}), '(1)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((42, 25, 42, 39), 'decimal.Decimal', 'Decimal', ({(42, 33, 42, 38): 'value'}, {}), '(value)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((42, 53, 42, 63), 'decimal.Decimal', 'Decimal', ({(42, 61, 42, 62): '(1)'}, {}), '(1)', False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((48, 29, 48, 47), 'math.floor', 'floor', ({(48, 35, 48, 46): '(value + 0.5)'}, {}), '(value + 0.5)', False, 'from math import ceil, floor, log2\n'), ((49, 21, 49, 38), 'math.ceil', 'ceil', ({(49, 26, 49, 37): '(value - 0.5)'}, {}), '(value - 0.5)', False, 'from math import ceil, floor, log2\n'), ((51, 15, 51, 26), 'math.ceil', 'ceil', ({(51, 20, 51, 25): 'value'}, {}), '(value)', False, 'from math import ceil, floor, log2\n')]