import warnings
from copy import deepcopy
from functools import singledispatch
from typing import Any, Iterable, Type, Union

import numpy as np

from ..parser.openqasm_ast import (
    ArrayLiteral,
    ArrayType,
    BitstringLiteral,
    BitType,
    BooleanLiteral,
    BoolType,
    ClassicalType,
    DiscreteSet,
    Expression,
    FloatLiteral,
    FloatType,
    Identifier,
    IndexedIdentifier,
    IntegerLiteral,
    IntType,
    UintType,
)

LiteralType = Union[BooleanLiteral, IntegerLiteral, FloatLiteral, ArrayLiteral, BitstringLiteral]


@singledispatch
def cast_to(into: Union[ClassicalType, Type[LiteralType]], variable: LiteralType) -> LiteralType:
    """Cast a variable into a given type. Order of parameters is to enable singledispatch"""
    if type(variable) == into:
        return variable
    if into == BooleanLiteral or isinstance(into, BoolType):
        return BooleanLiteral(bool(variable.value))
    if into == IntegerLiteral:
        return IntegerLiteral(int(variable.value))
    if into == FloatLiteral:
        return FloatLiteral(float(variable.value))
    raise TypeError(f"Cannot cast {type(variable).__name__} into {into.__name__}.")


@cast_to.register
def _(
    into: BitType, variable: Union[BooleanLiteral, ArrayLiteral, BitstringLiteral]
) -> ArrayLiteral:
    """
    Bit types can be sized or not, represented as Boolean literals or Array literals.
    Sized bit types can be instantiated with a Bitstring literal or Array literal.
    """
    if not into.size:
        return cast_to(BooleanLiteral, variable)
    size = into.size.value
    if isinstance(variable, BitstringLiteral):
        variable = convert_string_to_bool_array(variable)
    if (
        not all(isinstance(x, BooleanLiteral) for x in variable.values)
        or len(variable.values) != size
    ):
        raise ValueError(f"Invalid array to cast to bit register of size {size}: {variable}.")
    return ArrayLiteral(deepcopy(variable.values))


@cast_to.register
def _(into: IntType, variable: LiteralType) -> IntegerLiteral:
    """Cast to int with overflow warnings"""
    if isinstance(variable, ArrayLiteral):
        value = int("".join("01"[x.value] for x in variable.values[1:]), base=2)
        if variable.values[0].value:
            value *= -1
    else:
        value = variable.value
        if into.size is not None:
            limit = 2 ** (into.size.value - 1)
            value = int(np.sign(value) * (np.abs(int(value)) % limit))
            if value != variable.value:
                warnings.warn(
                    f"Integer overflow for value {variable.value} and size {into.size.value}."
                )
    return IntegerLiteral(value)


@cast_to.register
def _(into: UintType, variable: LiteralType) -> IntegerLiteral:
    """Cast to uint with overflow warnings. Bit registers can be cast to uint."""
    if isinstance(variable, ArrayLiteral):
        return IntegerLiteral(int("".join("01"[x.value] for x in variable.values), base=2))
    value = variable.value
    if into.size is not None:
        limit = 2**into.size.value
        value = int(value) % limit
        if value != variable.value:
            warnings.warn(
                f"Unsigned integer overflow for value {variable.value} and size {into.size.value}."
            )
    return IntegerLiteral(value)


@cast_to.register
def _(into: FloatType, variable: LiteralType) -> FloatLiteral:
    """Cast to float"""
    if into.size is None:
        value = float(variable.value)
    else:
        if into.size.value not in (16, 32, 64):
            raise ValueError("Float size must be one of {16, 32, 64}.")
        value = float(np.array(variable.value, dtype=np.dtype(f"float{into.size.value}")))
    return FloatLiteral(value)


@cast_to.register
def _(into: ArrayType, variable: Union[ArrayLiteral, DiscreteSet]) -> ArrayLiteral:
    """Cast to Array and enforce dimensions"""
    if len(variable.values) != into.dimensions[0].value:
        raise ValueError(
            f"Size mismatch between dimension of size {into.dimensions[0].value} "
            f"and values length {len(variable.values)}"
        )
    subtype = (
        ArrayType(into.base_type, into.dimensions[1:])
        if len(into.dimensions) > 1
        else into.base_type
    )
    return ArrayLiteral([cast_to(subtype, v) for v in variable.values])


def is_literal(expression: Expression) -> bool:
    return isinstance(
        expression,
        (
            BooleanLiteral,
            IntegerLiteral,
            FloatLiteral,
            BitstringLiteral,
            ArrayLiteral,
        ),
    )


def convert_string_to_bool_array(bit_string: BitstringLiteral) -> ArrayLiteral:
    """Convert BitstringLiteral to Boolean ArrayLiteral"""
    return ArrayLiteral(
        [BooleanLiteral(x == "1") for x in np.binary_repr(bit_string.value, bit_string.width)]
    )


def convert_bool_array_to_string(bit_string: ArrayLiteral) -> str:
    """Convert Boolean ArrayLiteral into a binary string"""
    return "".join(("1" if x.value else "0") for x in bit_string.values)


def is_none_like(value: Any) -> bool:
    """Returns whether value is None or an Array of Nones"""
    if isinstance(value, ArrayLiteral):
        return all(is_none_like(v) for v in value.values)
    return value is None


@singledispatch
def get_identifier_name(identifier: Union[Identifier, IndexedIdentifier]) -> str:
    """Get name of an identifier"""
    return identifier.name


@get_identifier_name.register
def _(identifier: IndexedIdentifier) -> str:
    """Get name of an indexed identifier"""
    return identifier.name.name


@singledispatch
def wrap_value_into_literal(value: Any) -> LiteralType:
    """Wrap a primitive variable into an AST node"""
    raise TypeError(f"Cannot wrap {value} into literal type")


@wrap_value_into_literal.register
def _(value: int) -> IntegerLiteral:
    return IntegerLiteral(value)


@wrap_value_into_literal.register
def _(value: float) -> FloatLiteral:
    return FloatLiteral(value)


@wrap_value_into_literal.register
def _(value: bool) -> BooleanLiteral:
    return BooleanLiteral(value)


@wrap_value_into_literal.register
def _(value: str) -> BitstringLiteral:
    return BitstringLiteral(value=int(value, base=2), width=len(value))


@wrap_value_into_literal.register(list)
def _(value: Iterable[Any]) -> ArrayLiteral:
    return ArrayLiteral([wrap_value_into_literal(v) for v in value])
