from functools import singledispatch
from jijmodeling.variables.element import Element
from jijmodeling.expression.constraint import Constraint, Penalty
from jijmodeling.expression.expression import Operator
from numbers import Number
from typing import Dict, Union, List, Tuple

from jijmodeling.variables.variable import Binary, DisNum, Placeholder
from jijmodeling.variables.array import ArraySizePlaceholder
from jijmodeling.variables.array import Tensor
from jijmodeling.expression.sum import Sum
from jijmodeling.transpilers.decode_from_sampleset import convert_to_int
from jijmodeling.transpilers.utils import _reshape_index
import numpy as np
import pyqubo as pyq

@singledispatch
def to_pyqubo(term, placeholder={}, fixed_variables={}, index_values={})->pyq.Base:
    # transpile to PyQUBO
    if isinstance(term, (Number, str)):
        return term
    raise TypeError("{} cannot convert to pyqubo object.".format(type(term)))

@to_pyqubo.register(Binary)
def binary_to_pyqubo(term, fixed_variables={}, relabel:str=None, **kwargs):
    if term.label in fixed_variables:
        return fixed_variables[term.label]
    else:
        # If to_pyqubo is called as an element of Tensor, 
        # use relabel to add subscript information to label.
        # ex. relabel => 'x[1][0]'
        term_label = term.label if relabel is None else relabel
        return pyq.Binary(term_label)


@to_pyqubo.register(DisNum)
def disnum_to_pyqubo(term: DisNum, fixed_variables={}, relabel: str=None, placeholder={}, **kwargs):
    if term.label in fixed_variables:
        return fixed_variables[term.label]
    var_label = (term.label if relabel is None else relabel) + '[{}]'
    upper: int = convert_to_int(term.upper, placeholder=placeholder)
    lower: int = convert_to_int(term.lower, placeholder=placeholder)
    bits: int = convert_to_int(term.bits, placeholder=placeholder)
    coeff = (upper - lower)/(2**bits - 1)
    return coeff * sum(2**i * pyq.Binary(var_label.format(i)) for i in range(bits)) + lower

@to_pyqubo.register(Placeholder)
def placeholder_to_pyqubo(term, placeholder={}, index_values={}, fixed_variables={}):
    if term.label in placeholder:
        return placeholder[term.label]
    else:
        return pyq.Placeholder(term.label)

@to_pyqubo.register(Element)
def element_to_pyqubo(term, placeholder={}, index_values={}, fixed_variables={}):
    if term.array_label is None:
        start, end = term.range
        se_to_pyq = lambda a: to_pyqubo(a, placeholder=placeholder, index_values=index_values, fixed_variables=fixed_variables)
        data_set = np.arange(se_to_pyq(start), se_to_pyq(end))
    else:
        data_set = placeholder[term.array_label]
    return data_set[index_values[term.label]]

@to_pyqubo.register(ArraySizePlaceholder)
def arraysizeplaceholder_to_pyqubo(term: ArraySizePlaceholder, placeholder={}, index_values={}, fixed_variables={}):
    if term.array_label in placeholder:
        p_array = np.array(placeholder[term.array_label])
        return p_array.shape[term.dimension]
    else:
        raise ValueError("The placeholder must be set to list or np.ndarray if the size of the array is used.")


@to_pyqubo.register(Tensor)
def tensor_to_pyqubo(
    term: Tensor, 
    index_values: dict, 
    placeholder:Dict[str, Union[Number, list, np.ndarray]]={}, 
    fixed_variables:Dict[str, Dict[tuple, Number]]={}):

    ind_value_list = [index_values[ind] if isinstance(ind, str) else to_pyqubo(ind, placeholder=placeholder, index_values=index_values) 
                      for ind in term.index_labels]

    # If a fixed value is set, it is returned.
    if term.label in fixed_variables:
        if tuple(ind_value_list) in fixed_variables[term.label]:
            fixed_var: Union[dict, Number] = fixed_variables[term.label]
            if isinstance(fixed_var, dict):
                return fixed_var[tuple(ind_value_list)]
            else:
                raise TypeError("fixed variable's value type is dict")
    if term.label in placeholder:
        if isinstance(placeholder[term.label], (list, np.ndarray)):
            return np.array(placeholder[term.label])[tuple(ind_value_list)]
    
    # make name of tensor element (ex. x[0][1])
    tensor_name = term.label + ''.join(['[{}]'.format(ind) for ind in ind_value_list])
    pyq_obj = to_pyqubo(term.variable, relabel=tensor_name)
    return pyq_obj


@to_pyqubo.register(Operator)
def operator_to_pyqubo(term: Operator, placeholder={}, fixed_variables={}, index_values={}):
    pyq_children = [to_pyqubo(c, 
                              placeholder=placeholder, 
                              fixed_variables=fixed_variables, 
                              index_values=index_values) 
                    for c in term.children]
    return term.operation(pyq_children)


@to_pyqubo.register(Sum)
def sum_to_pyqubo(term: Sum, placeholder={}, fixed_variables={}, index_values={}):

    # convert index set to list
    ind_value_list = _reshape_index(term.indices, index_values, placeholder)

    return sum(to_pyqubo(term.children[0], placeholder=placeholder, fixed_variables=fixed_variables, index_values=ind) for ind in ind_value_list)


@to_pyqubo.register(Constraint)
def constraint_to_pyqubo(term: Constraint, placeholder={}, fixed_variables={}, index_values={}):
    if term._penalty is None:
        return 0

    pena = to_pyqubo(term.penalty, placeholder=placeholder, fixed_variables=fixed_variables, index_values=index_values)

    return pyq.Constraint(pena, term.label)

@to_pyqubo.register(Penalty)
def penalty_to_pyqubo(term: Penalty, placeholder={}, fixed_variables={}, index_values={}):
    pena_term = term.children[0]
    return to_pyqubo(pena_term, placeholder=placeholder, fixed_variables=fixed_variables, index_values=index_values)
