# Copyright 2021 Jij Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

import openjij
from .model import make_BinaryQuadraticModel, make_BinaryQuadraticModel_from_JSON, BinaryQuadraticModel


def make_KingGraph(linear=None, quadratic=None, king_graph=None):
    """ KingGraph factory
    Returns:
        generated KingGraph class
    """
    mock_linear = {}
    mock_quadratic = {}

    if linear != None:
        mock_linear = linear

    if quadratic != None:
        mock_quadratic = quadratic

    if mock_linear == {} and mock_quadratic == {}:
        # no elements in linear and quadratic
        # fetch first element of the king_graph
        label = (king_graph[0][0], king_graph[0][0])
        # add to linear
        mock_linear[label] = 1.0

    class KingGraph(make_BinaryQuadraticModel(mock_linear, mock_quadratic)):
        """
        BQM for king graph of HITACHI CMOS Annealer
        Attributes:
            xrange (list(int)): represents hardware (CMOS) restricts for coordinate. [xmin, xmax]
            yrange (list(int)): represents hardware (CMOS) restricts for coordinate. [ymin, ymax]
            prange (list(int)): represents hardware (CMOS) restricts for the strength of interactions 'p'. [pmin, pmax]
            king_graph (list(int)): 
                Annealing cloud Web API format representation of interaction coefficients
                Quadratic term [x1, y1, x2, y2, value]
                Linear term    [x1, y1, x1, y1, value]
        """
    
        def __init__(self, linear: dict=None, quadratic: dict=None, offset: float=0.0,

                king_graph=None, vartype=openjij.SPIN, machine_type: str=''):
            """__init__.

            Args:
                linear (dict): linear biases
                quadratic (dict): quadratic biases
                offset (float): offset
                king_graph: represents ising or QUBO interaction.
                    Each spins are decided by 2-d corrdinate x, y.
                    Quadratic term [x1, y1, x2, y2, value]
                    Linear term    [x1, y1, x1, y1, value]
                vartype: 'SPIN' or 'BINARY'
                machine_type (str): choose 'ASIC' or 'FPGA'
            """
            vartype = openjij.cast_vartype(vartype)
    
            # set parameter ranges
            self.machine_type = machine_type
            if self.machine_type == "ASIC":
                self.xrange = [0, 351+1]
                self.yrange = [0, 175+1]
                self.prange = [-3, 3]
            elif self.machine_type == "FPGA":
                self.xrange = [0, 79+1]
                self.yrange = [0, 79+1]
                self.prange = [-127, 127]
            else:
                raise ValueError('machine type should be ASIC or FPGA')
    
            # convert format h, J, Q and initilize BQM
            if king_graph is not None:
                linear, quadratic = self._convert_to_BQM_format(
                    king_graph, vartype)
            super().__init__(linear, quadratic, offset=offset, vartype=vartype)
    
            # reformat to ising king graph (which is Web API format)
            if king_graph is not None and vartype == openjij.SPIN:
                self._ising_king_graph = king_graph
    
            else:
                # generate Ising h and J and create ising_king_graph format
                lin, quad, _ = self.to_ising()
                self._ising_king_graph = []
                for index, h in lin.items():
                    if h != 0:
                        x, y = self._convert_to_xy(index)
                        self._ising_king_graph.append([x, y, x, y, h])
                for (i, j), J in quad.items():
                    if J != 0:
                        x1, y1 = self._convert_to_xy(i)
                        x2, y2 = self._convert_to_xy(j)
                        self._ising_king_graph.append([x1, y1, x2, y2, J])
    
            self._validation_ising_king_graph()
    
        def _convert_to_BQM_format(self, king_graph, vartype):
            linear, quad = {}, {}
            for x1, y1, x2, y2, value in king_graph:
                if (x1, y1) == (x2, y2):
                    linear[(x1, y1)] = value
                else:
                    quad[(x1, y1), (x2, y2)] = value
            return linear, quad
    
        def get_ising_king_graph(self):
            return self._ising_king_graph
    
        def king_indices(self):
            if isinstance(self.indices[0], tuple):
                return self.indices
            else:
                return [self._convert_to_xy(i) for i in self.indices]
    
        def _convert_to_xy(self, index):
            if isinstance(index, tuple):
                return index[0], index[1]
            else:
                y = int(index / self.xrange[1])
                return int(index - y * self.xrange[1]), y
    
        def convert_to_index(self, x, y):
            return y * self.xrange[1] + x
    
        def _validation_ising_king_graph(self):
            for xi, yi, xj, yj, p in self._ising_king_graph:
                if yi >= self.yrange[1] or yj >= self.yrange[1]:
                    raise ValueError('Graph is incomplete xi: {}, yi: {}, xj: {}, yj: {}, p:{}'
                                     .format(xi, yi, xj, yj, p))
                if not (xi in [xj, xj-1, xj+1]) or not (yi in [yj, yj-1, yj+1]):
                    raise ValueError('Graph is incomplete xi: {}, yi: {}, xj: {}, yj: {}, p:{}'
                                     .format(xi, yi, xj, yj, p))
                if not (self.prange[0] <= p <= self.prange[1]):
                    raise ValueError('Graph is incomplete xi: {}, yi: {}, xj: {}, yj: {}, p: {}'
                                     .format(xi, yi, xj, yj, p))
    
    
        def energy(self, sample):
            return super().energy(sample, sparse=True)
    
        def energies(self, samples_like):
            return super().energies(sample_like, sparse=True)

    return KingGraph

def make_KingGraph_from_JSON(obj):
    """ KingGraph factory for JSON
    Args:
        obj (dict): JSON object
    Returns:
        generated KingGraph class
    """
    label = obj['variable_labels'][0]
    if isinstance(label, list):
        #convert to tuple
        label = tuple(label)

    mock_linear = {label:1.0}

    return make_KingGraph(mock_linear, {})

def KingGraph(linear=None, quadratic=None, offset=0.0,
                     king_graph=None, vartype=openjij.SPIN, machine_type=''):
    """generate KingGraph model.

    Args:
        linear (dict): linear biases
        quadratic (dict): quadratic biases
        offset (float): offset
        king_graph: represents ising or QUBO interaction.
            Each spins are decided by 2-d corrdinate x, y.
            
            * Quadratic term: [x1, y1, x2, y2, value]
            * Linear term:    [x1, y1, x1, y1, value]
        vartype: 'SPIN' or 'BINARY'
        machine_type (str): choose 'ASIC' or 'FPGA'
    Returns:
        generated KingGraphModel
    Examples:
        The following code shows intialization of KingGraph::

            >>> h = {}
            >>> J = {(0, 1): -1.0, (1, 2): -3.0}
            >>> king_graph = oj.KingGraph(machine_type="ASIC", linear=h, quadratic=J)

        You can initialize it from `king_interaction`::

            >>> king_interaction = [[0, 0, 1, 0, -1.0], [1, 0, 2, 0, -3.0]]
            >>> king_graph = oj.KingGraph(machine_type="ASIC", king_graph=king_interaction)
    """

    Model = make_KingGraph(linear, quadratic, king_graph)

    return Model(linear, quadratic, offset, king_graph, vartype, machine_type)

#classmethods
KingGraph.from_qubo = \
        lambda Q, offset=0.0, **kwargs: make_KingGraph({}, Q).from_qubo(Q, offset, **kwargs)

KingGraph.from_ising = \
        lambda linear, quadratic, offset=0.0, **kwargs: make_KingGraph(linear, quadratic).from_ising(linear, quadratic, offset, **kwargs)

KingGraph.from_serializable = \
        lambda obj: make_KingGraph_from_JSON(obj).from_serializable(obj)

