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)))

[((1397, 1411), 'decimal.Decimal', 'Decimal', (['value'], {}), '(value)\n', (1404, 1411), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1425, 1435), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (1432, 1435), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((3472, 3489), 'torch.sign', 'torch.sign', (['value'], {}), '(value)\n', (3482, 3489), False, 'import torch\n'), ((4434, 4452), 'math.log2', 'log2', (['(sign * value)'], {}), '(sign * value)\n', (4438, 4452), False, 'from math import ceil, floor, log2\n'), ((3597, 3614), 'torch.sign', 'torch.sign', (['value'], {}), '(value)\n', (3607, 3614), False, 'import torch\n'), ((3714, 3737), 'torch.ceil', 'torch.ceil', (['(value - 0.5)'], {}), '(value - 0.5)\n', (3724, 3737), False, 'import torch\n'), ((1546, 1560), 'decimal.Decimal', 'Decimal', (['value'], {}), '(value)\n', (1553, 1560), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1574, 1584), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (1581, 1584), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1636, 1650), 'decimal.Decimal', 'Decimal', (['value'], {}), '(value)\n', (1643, 1650), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1664, 1674), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (1671, 1674), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((3802, 3826), 'torch.floor', 'torch.floor', (['(value + 0.5)'], {}), '(value + 0.5)\n', (3813, 3826), False, 'import torch\n'), ((1787, 1801), 'decimal.Decimal', 'Decimal', (['value'], {}), '(value)\n', (1794, 1801), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1815, 1825), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (1822, 1825), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1879, 1893), 'decimal.Decimal', 'Decimal', (['value'], {}), '(value)\n', (1886, 1893), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((1907, 1917), 'decimal.Decimal', 'Decimal', (['(1)'], {}), '(1)\n', (1914, 1917), False, 'from decimal import ROUND_HALF_DOWN, ROUND_HALF_EVEN, ROUND_HALF_UP, Decimal\n'), ((2291, 2309), 'math.floor', 'floor', (['(value + 0.5)'], {}), '(value + 0.5)\n', (2296, 2309), False, 'from math import ceil, floor, log2\n'), ((2331, 2348), 'math.ceil', 'ceil', (['(value - 0.5)'], {}), '(value - 0.5)\n', (2335, 2348), False, 'from math import ceil, floor, log2\n'), ((2408, 2419), 'math.ceil', 'ceil', (['value'], {}), '(value)\n', (2412, 2419), False, 'from math import ceil, floor, log2\n')]