# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.

"""AutoQASM types and type utilities."""

from __future__ import annotations

from collections.abc import Iterable
from typing import Any, List, Union, get_args

import numpy as np
import oqpy
import oqpy.base
from braket.circuits import FreeParameterExpression
from braket.registers import Qubit
from openpulse import ast

from autoqasm import errors, program


def is_qasm_type(val: Any) -> bool:
    """Returns whether the provided object is of a QASM type or is itself a QASM type
    which receives special treatment by AutoQASM.

    Args:
        val (Any): The object of which to check the type.

    Returns:
        bool: Whether the object is a QASM type.
    """
    qasm_types = (
        oqpy.Range,
        oqpy._ClassicalVar,
        oqpy.Qubit,
        oqpy.base.OQPyExpression,
        FreeParameterExpression,
    )
    # The input can either be a class, like oqpy.Range ...
    if type(val) is type:
        return issubclass(val, qasm_types)
    # ... or an instance of a class, like oqpy.Range(10)
    return isinstance(val, qasm_types)


def make_annotations_list(annotations: str | Iterable[str] | None) -> List[str]:
    return [annotations] if isinstance(annotations, str) else annotations or []


QubitIdentifierType = Union[
    int, str, Qubit, oqpy._ClassicalVar, oqpy.base.OQPyExpression, oqpy.Qubit
]


def is_qubit_identifier_type(qubit: Any) -> bool:
    """Checks if a given object is a qubit identifier type.

    Args:
        qubit (Any): The object to check.

    Returns:
        bool: True if the object is a qubit identifier type, False otherwise.
    """
    return isinstance(qubit, get_args(QubitIdentifierType))


class Range(oqpy.Range):
    def __init__(
        self,
        start: int,
        stop: int | None = None,
        step: int | None = 1,
        annotations: str | Iterable[str] | None = None,
    ):
        """Creates a range definition.

        Args:
            start (int): Start of the range.
            stop (int | None): End of the range. Defaults to None.
            step (int | None): Step of the range. Defaults to 1.
            annotations (str | Iterable[str] | None): Annotations for the range.
        """
        if stop is None:
            stop = start
            start = 0
        super(Range, self).__init__(start, stop, step)
        self.annotations = make_annotations_list(annotations)


class ArrayVar(oqpy.ArrayVar):
    def __init__(
        self,
        init_expression: Iterable,
        *args,
        annotations: str | Iterable[str] | None = None,
        **kwargs,
    ):
        if (
            program.get_program_conversion_context().subroutines_processing
            or not program.get_program_conversion_context().at_function_root_scope
        ):
            raise errors.InvalidArrayDeclaration(
                "Arrays may only be declared at the root scope of an AutoQASM main function."
            )

        if not isinstance(init_expression, Iterable):
            raise errors.InvalidArrayDeclaration("init_expression must be an iterable type.")

        dimensions = np.shape(init_expression)
        super(ArrayVar, self).__init__(
            init_expression=init_expression,
            *args,
            annotations=make_annotations_list(annotations),
            dimensions=dimensions,
            **kwargs,
        )
        self.name = program.get_program_conversion_context().next_var_name(oqpy.ArrayVar)


class BitVar(oqpy.BitVar):
    def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
        super(BitVar, self).__init__(
            *args, annotations=make_annotations_list(annotations), **kwargs
        )
        self.name = program.get_program_conversion_context().next_var_name(oqpy.BitVar)
        if self.size and self.init_expression != "output":
            value = self.init_expression or 0
            self.init_expression = ast.BitstringLiteral(value, self.size)


class BoolVar(oqpy.BoolVar):
    def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
        super(BoolVar, self).__init__(
            *args, annotations=make_annotations_list(annotations), **kwargs
        )
        self.name = program.get_program_conversion_context().next_var_name(oqpy.BoolVar)


class FloatVar(oqpy.FloatVar):
    def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
        super(FloatVar, self).__init__(
            *args, annotations=make_annotations_list(annotations), **kwargs
        )
        self.name = program.get_program_conversion_context().next_var_name(oqpy.FloatVar)


class IntVar(oqpy.IntVar):
    def __init__(self, *args, annotations: str | Iterable[str] | None = None, **kwargs):
        super(IntVar, self).__init__(
            *args, annotations=make_annotations_list(annotations), **kwargs
        )
        self.name = program.get_program_conversion_context().next_var_name(oqpy.IntVar)
