# PyZX - Python library for quantum circuit rewriting
#        and optimization using the ZX-calculus
# Copyright (C) 2018 - Aleks Kissinger and John van de Wetering

# 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.

"""
NEW VERSION 
"""

import json
from collections import defaultdict
from fractions import Fraction

from typing import Callable, Optional, List, Dict, Tuple

from pyzx_param.utils import EdgeType, VertexType, FractionLike, toggle_edge, vertex_is_zx, toggle_vertex
from pyzx_param.graph.base import BaseGraph, VT, ET, upair
import pyzx_param.rewrite_rules.rules as rules
import pyzx_param.rewrite_rules.hrules as hrules


def match_X_spiders(
        g: BaseGraph[VT, ET],
        vertexf: Optional[Callable[[VT], bool]] = None
        ) -> List[VT]:
    if vertexf is not None: candidates = set([v for v in g.vertices() if vertexf(v)])
    else: candidates = g.vertex_set()
    types = g.types()

    return [v for v in candidates if types[v] == VertexType.X]


def match_Z_spiders(
        g: BaseGraph[VT, ET],
        vertexf: Optional[Callable[[VT], bool]] = None
        ) -> List[VT]:
    if vertexf is not None: candidates = set([v for v in g.vertices() if vertexf(v)])
    else: candidates = g.vertex_set()
    types = g.types()

    return [v for v in candidates if types[v] == VertexType.Z]


def color_change(g: BaseGraph[VT,ET], matches: List[VT]) -> rules.RewriteOutputType[VT,ET]:
    for v in matches:
        g.set_type(v, toggle_vertex(g.type(v)))
        for e in g.incident_edges(v):
            g.set_edge_type(e, toggle_edge(g.edge_type(e)))
    return ({}, [],[],False)


def pauli_matcher(
        g: BaseGraph[VT,ET], 
        vertexf: Optional[Callable[[VT],bool]] = None
        ) -> List[Tuple[VT,VT]]:
    if vertexf is not None: candidates = set([v for v in g.vertices() if vertexf(v)])
    else: candidates = g.vertex_set()
    phases = g.phases()
    types = g.types()
    m: List[Tuple[VT,VT]] = []
    paulis = {v for v in candidates
                if phases[v] == 1 and vertex_is_zx(types[v])}
    if not paulis: return m
    while len(candidates) > 0:
        v = candidates.pop()
        if v in paulis and g.vertex_degree(v) == 2: continue
        for w in g.neighbors(v):
            if w in paulis: break
        else:
            continue
        et = g.edge_type(g.edge(v,w))
        if ((types[v] == types[w] and et == EdgeType.HADAMARD) or
            (vertex_is_zx(types[v]) and types[v] != types[w] and et == EdgeType.SIMPLE) or
            (types[v] == VertexType.H_BOX and phases[v] == 1 and (
                (et == EdgeType.SIMPLE and types[w] == VertexType.X) or
                (et == EdgeType.HADAMARD and types[w] == VertexType.Z)))
            ):
            m.append((w,v))
            candidates.difference_update(g.neighbors(v))
            candidates.difference_update(g.neighbors(w))
    return m


def pauli_push(g: BaseGraph[VT,ET],
               matches: List[Tuple[VT,VT]]
               ) -> rules.RewriteOutputType[VT,ET]:
    """Pushes a Pauli (i.e. a pi phase) through another spider."""
    rem_verts: List[VT] = []
    rem_edges: List[ET] = []
    etab: Dict[Tuple[VT,VT], List[int]] = dict()
    for w,v in matches:  # w is a Pauli and v is the spider we are gonna push it through
        if g.vertex_degree(w) == 2:
            rem_verts.append(w)
            l = list(g.neighbors(w))
            l.remove(v)
            v2 = l[0]
            et1 = g.edge_type(g.edge(v,w))
            et2 = g.edge_type(g.edge(v2,w))
            etab[upair(v,v2)] = [1,0] if et1 == et2 else [0,1]
        else:
            g.set_phase(w,0)

        new_verts = []
        if vertex_is_zx(g.type(v)):
            g.scalar.add_phase(g.phase(v))
            g.set_phase(v,(-g.phase(v)) % 2)
            t = toggle_vertex(g.type(v))
            p: FractionLike = Fraction(1)
        else:
            t = VertexType.Z
            p = 0
        for edge in g.incident_edges(v):
            st = g.edge_st(edge)
            n = st[0] if st[1] == v else st[1]
            if n == w: continue
            r = 0.5*(g.row(n) + g.row(v))
            q = 0.5*(g.qubit(n) + g.qubit(v))
            et = g.edge_type(edge)
            rem_edges.append(edge)
            w2 = g.add_vertex(t,q,r,p)
            etab[upair(v,w2)] = [1,0]
            etab[upair(n,w2)] = [1,0] if et == EdgeType.SIMPLE else [0,1]
            new_verts.append(w2)
        if not vertex_is_zx(g.type(v)): # v is H_BOX
            if len(new_verts) == 2:
                etab[upair(new_verts[0],new_verts[1])] = [0,1]
            else:
                r = (g.row(v) + sum(g.row(n) for n in new_verts)) / (len(new_verts) + 1)
                q = (g.qubit(v) + sum(g.qubit(n) for n in new_verts))/(len(new_verts)+1)
                h = g.add_vertex(VertexType.H_BOX,q,r,Fraction(1))
                for n in new_verts: etab[upair(h,n)] = [1,0]
    return (etab, rem_verts, rem_edges, False)


def match_hadamard_edge(
        g: BaseGraph[VT,ET], 
        edgef: Optional[Callable[[ET],bool]] = None
        ) -> List[ET]:
    if edgef is not None: candidates = set([e for e in g.edges() if edgef(e)])
    else: candidates = g.edge_set()
    return [e for e in candidates if g.edge_type(e)==EdgeType.HADAMARD]


def match_edge(
        g: BaseGraph[VT,ET], 
        edgef: Optional[Callable[[ET],bool]] = None
        ) -> List[ET]:
    if edgef is not None: candidates = set([e for e in g.edges() if edgef(e)])
    else: candidates = g.edge_set()
    return list(candidates)


def euler_expansion(g: BaseGraph[VT,ET], 
                    matches: List[ET]
                    ) -> rules.RewriteOutputType[VT,ET]:
    """Expands the given Hadamard-edges into pi/2 phases using its Euler decomposition."""
    types = g.types()
    phases = g.phases()
    rem_edges = []
    etab = {}
    for e in matches:
        rem_edges.append(e)
        v1,v2 = g.edge_st(e)
        if vertex_is_zx(types[v1]) and types[v1] == types[v2]:
            r = 0.5*(g.row(v1) + g.row(v2))
            q = 0.5*(g.qubit(v1) + g.qubit(v2))
            t = toggle_vertex(types[v1])
            v = g.add_vertex(t,q,r)
            etab[upair(v,v1)] = [1,0]
            etab[upair(v,v2)] = [1,0]
            if phases[v1] == Fraction(1,2) or phases[v2] == Fraction(1,2):
                g.add_to_phase(v1,Fraction(3,2))
                g.add_to_phase(v2,Fraction(3,2))
                g.set_phase(v, Fraction(3,2))
                g.scalar.add_phase(Fraction(1,4))
            else:
                g.add_to_phase(v1,Fraction(1,2))
                g.add_to_phase(v2,Fraction(1,2))
                g.set_phase(v, Fraction(1,2))
                g.scalar.add_phase(Fraction(7,4))
        else:
            r = 0.25*g.row(v1) + 0.75*g.row(v2)
            q = 0.25*g.qubit(v1) + 0.75*g.qubit(v2)
            w1 = g.add_vertex(VertexType.Z,q,r,Fraction(1,2))
            etab[upair(v2,w1)] = [1,0]
            r = 0.5*g.row(v1) + 0.5*g.row(v2)
            q = 0.5*g.qubit(v1) + 0.5*g.qubit(v2)
            w2 = g.add_vertex(VertexType.X,q,r,Fraction(1,2))
            etab[upair(w1,w2)] = [1,0]
            r = 0.75*g.row(v1) + 0.25*g.row(v2)
            q = 0.75*g.qubit(v1) + 0.25*g.qubit(v2)
            w3 = g.add_vertex(VertexType.Z,q,r,Fraction(1,2))
            etab[upair(w2,w3)] = [1,0]
            etab[upair(w3,v1)] = [1,0]
            g.scalar.add_phase(Fraction(7,4))
            
    return (etab, [], rem_edges, False)

def add_Z_identity(g: BaseGraph[VT,ET], 
        matches: List[ET]
        ) -> rules.RewriteOutputType[VT,ET]:
    rem_edges = []
    etab = {}
    for e in matches:
        rem_edges.append(e)
        et = g.edge_type(e)
        v1,v2 = g.edge_st(e)
        r = 0.5*(g.row(v1) + g.row(v2))
        q = 0.5*(g.qubit(v1) + g.qubit(v2))
        w = g.add_vertex(VertexType.Z, q,r, 0)
        etab[upair(v1,w)] = [1,0] if et == EdgeType.SIMPLE else [0,1]
        etab[upair(v2,w)] = [1,0]
    return (etab, [], rem_edges, False)

def match_bialgebra(g: BaseGraph[VT,ET], 
        edgef: Optional[Callable[[ET],bool]] = None
        ) -> List[Tuple[VT,VT]]:
    if edgef is not None: candidates = set([e for e in g.edges() if edgef(e)])
    else: candidates = g.edge_set()
    m = []
    types = g.types()
    phases = g.phases()
    while len(candidates) > 0:
        e = candidates.pop()
        if g.edge_type(e) != EdgeType.SIMPLE: continue
        v,w = g.edge_st(e)
        if types[v] != VertexType.X:
            v,w = w,v
        if types[v] != VertexType.X: continue
        if types[w] == VertexType.Z:
            if phases[v] != 0 or phases[w] != 0: continue
            m.append((v,w))
            for n in g.neighbors(v):
                candidates.difference_update(g.incident_edges(n))
            for n in g.neighbors(w):
                candidates.difference_update(g.incident_edges(n))
        if types[w] == VertexType.H_BOX:
            if phases[v] != 0 or phases[w] != 1: continue
            m.append((v,w))
            for n in g.neighbors(v):
                candidates.difference_update(g.incident_edges(n))
            for n in g.neighbors(w):
                candidates.difference_update(g.incident_edges(n))
    return m

def bialgebra(g: BaseGraph[VT,ET],
        matches: List[Tuple[VT,VT]]
        ) -> rules.RewriteOutputType[VT,ET]:
    rem_verts = []
    etab = {}
    for v1, v2 in matches:
        rem_verts.append(v1)
        rem_verts.append(v2)
        v = (v1,v2)
        new_verts: Tuple[List[VT],List[VT]] = ([],[]) # new vertices for v1 and v2

        for i, j in [(0, 1), (1, 0)]:
            multi_edge_found = False
            for e in g.incident_edges(v[i]):
                source, target = g.edge_st(e)
                other_vertex = source if source != v[i] else target
                if other_vertex != v[j] or multi_edge_found:
                    q = 0.4*g.qubit(other_vertex) + 0.6*g.qubit(v[i])
                    r = 0.4*g.row(other_vertex) + 0.6*g.row(v[i])
                    newv = g.add_vertex(g.type(v[j]), qubit=q, row=r)
                    g.set_phase(newv, g.phase(v[j]))
                    new_verts[i].append(newv)
                    if other_vertex == v[j]:
                        q = 0.4*g.qubit(v[i]) + 0.6*g.qubit(other_vertex)
                        r = 0.4*g.row(v[i]) + 0.6*g.row(other_vertex)
                        newv2 = g.add_vertex(g.type(v[i]), qubit=q, row=r)
                        new_verts[j].append(newv2)
                        other_vertex = newv2
                    if upair(newv, other_vertex) not in etab:
                        etab[upair(newv, other_vertex)] = [0, 0]
                    type_index = 0 if g.edge_type(e) == EdgeType.SIMPLE else 1
                    etab[upair(newv, other_vertex)][type_index] += 1
                elif i == 0: # only add new vertex once
                    multi_edge_found = True

        for n1 in new_verts[0]:
            for n2 in new_verts[1]:
                if upair(n1,n2) not in etab:
                    etab[upair(n1,n2)] = [0, 0]
                etab[upair(n1,n2)][0] += 1

        if g.type(v1) == VertexType.H_BOX or g.type(v2) == VertexType.H_BOX: # x-h bialgebra
            x_vertex = v1 if g.type(v2) == VertexType.H_BOX else v2
            g.scalar.add_power(g.vertex_degree(x_vertex)-2)
        else: # z-x bialgebra
            g.scalar.add_power((g.vertex_degree(v1)-2)*(g.vertex_degree(v2)-2))
    return (etab, rem_verts, [], False)

def match_bialgebra_op(g: BaseGraph[VT,ET],
        vertexf: Optional[Callable[[VT], bool]] = None,
        vertex_type: Optional[Tuple[VertexType, VertexType]] = None,
        edge_type: Optional[EdgeType] = None
        ) -> Optional[Tuple[List[VT], List[VT]]]:
    if vertexf is not None: candidates = set([v for v in g.vertices() if vertexf(v)])
    else: candidates = g.vertex_set()
    if vertex_type is not None:
        vtype1, vtype2 = vertex_type
    else:
        vtype1, vtype2 = VertexType.Z, VertexType.X
    if edge_type is None:
        edge_type = EdgeType.SIMPLE
    type1_vertices = [v for v in candidates if g.type(v) == vtype1]
    type2_vertices = [v for v in candidates if g.type(v) == vtype2]
    if len(type1_vertices) <= 1 or len(type2_vertices) <= 1:
        return None
    #the vertices must be phase-free
    for v in type1_vertices + type2_vertices:
        if g.phase(v) != 0:
            return None
    # the vertices must have one external edge
    for v1 in type1_vertices:
        if g.vertex_degree(v1) != len(type2_vertices) + 1:
            return None
    for v2 in type2_vertices:
        if g.vertex_degree(v2) != len(type1_vertices) + 1:
            return None
    # if all type1 vertices are connected to all type2 vertices with a simple edge, then they are a match
    for v1 in type1_vertices:
        for v2 in type2_vertices:
            edges = list(g.edges(v1, v2))
            if not (len(edges) == 1 and g.edge_type(edges[0]) == edge_type):
                return None
    return type1_vertices, type2_vertices

def bialgebra_op(g: BaseGraph[VT,ET],
        matches: Tuple[List[VT], List[VT]],
        edge_type: Optional[EdgeType] = EdgeType.SIMPLE
        ) -> rules.RewriteOutputType[VT,ET]:
    def get_neighbors_and_loops(type1_vertices: List[VT], type2_vertices: List[VT]) -> Tuple[List[Tuple[VT, EdgeType]], List[EdgeType]]:
        neighbors: List[Tuple[VT, EdgeType]] = []
        loops: List[EdgeType] = []
        for v1 in type1_vertices:
            for edge in g.incident_edges(v1):
                edge_st = g.edge_st(edge)
                neighbor = edge_st[0] if edge_st[0] != v1 else edge_st[1]
                if neighbor in type2_vertices:
                    continue
                elif neighbor in type1_vertices:
                    if v1 > neighbor:
                        loops.append(g.edge_type(edge))
                else:
                    neighbors.append((neighbor, g.edge_type(edge)))
        return neighbors, loops

    def add_vertex_with_averages(vertices, g, vtype):
        average_row = sum(g.row(v) for v in vertices) / len(vertices)
        average_qubit = sum(g.qubit(v) for v in vertices) / len(vertices)
        return g.add_vertex(vtype, average_qubit, average_row)

    def update_etab(etab, new_vertex, neighbors, loops):
        for n, et in neighbors + [(new_vertex, et) for et in loops]:
            etab[upair(new_vertex, n)][0 if et == EdgeType.SIMPLE else 1] += 1

    type1_vertices, type2_vertices = matches
    neighbors1, loops1 = get_neighbors_and_loops(type1_vertices, type2_vertices)
    neighbors2, loops2 = get_neighbors_and_loops(type2_vertices, type1_vertices)

    new_vertex1 = add_vertex_with_averages(type1_vertices, g, g.type(type2_vertices[0]))
    new_vertex2 = add_vertex_with_averages(type2_vertices, g, g.type(type1_vertices[0]))

    etab: dict = defaultdict(lambda: [0, 0])
    if edge_type == EdgeType.SIMPLE:
        etab[upair(new_vertex1, new_vertex2)] = [1, 0]
    else:
        etab[upair(new_vertex1, new_vertex2)] = [0, 1]
    update_etab(etab, new_vertex1, neighbors1, loops1)
    update_etab(etab, new_vertex2, neighbors2, loops2)

    g.scalar.add_power(-(len(neighbors1)-1)*(len(neighbors2)-1))

    return (etab, type1_vertices + type2_vertices, [], False)


MATCHES_VERTICES = 1
MATCHES_EDGES = 2

operations = {
    "spider": {"text": "fuse spiders", 
               "tooltip": "Fuses connected spiders of the same color",
               "matcher": rules.match_spider_parallel, 
               "rule": rules.spider, 
               "type": MATCHES_VERTICES},
    "to_z": {"text": "change color to Z", 
               "tooltip": "Changes X spiders into Z spiders by pushing out Hadamards",
               "matcher": match_X_spiders, 
               "rule": color_change, 
               "type": MATCHES_VERTICES},
    "to_x": {"text": "change color to X", 
               "tooltip": "Changes Z spiders into X spiders by pushing out Hadamards",
               "matcher": match_Z_spiders, 
               "rule": color_change, 
               "type": MATCHES_VERTICES},
    "rem_id": {"text": "remove identity", 
               "tooltip": "Removes a 2-ary phaseless spider",
               "matcher": rules.match_ids_parallel, 
               "rule": rules.remove_ids, 
               "type": MATCHES_VERTICES},
    "id_z": {"text": "Add Z identity", 
               "tooltip": "Adds a phaseless arity 2 Z spider on the selected edges",
               "matcher": match_edge, 
               "rule": add_Z_identity, 
               "type": MATCHES_EDGES},
    "z_to_z_box": {"text": "Convert Z-spider to Z-box",
                "tooltip": "Converts a Z-spider into a Z-box",
                "matcher": rules.match_z_to_z_box_parallel,
                "rule": rules.z_to_z_box,
                "type": MATCHES_VERTICES},
    "hopf": {"text": "Hopf rule",
             "tooltip": "Applies the Hopf rule between pairs of spiders that share parallel edges",
             "matcher": rules.match_hopf,
             "rule": rules.hopf,
             "type": MATCHES_VERTICES},
    "remove_self_loops": {"text": "remove self-loops",
                        "tooltip": "Removes all self-loops on a spider",
                        "matcher": rules.match_self_loop,
                        "rule": rules.remove_self_loops,
                        "type": MATCHES_VERTICES},
    "had2edge": {"text": "Convert H-box", 
               "tooltip": "Converts an arity 2 H-box into an H-edge.",
               "matcher": hrules.match_hadamards, 
               "rule": hrules.hadamard_to_h_edge, 
               "type": MATCHES_VERTICES},
    "fuse_hbox": {"text": "Fuse H-boxes", 
               "tooltip": "Merges two adjacent H-boxes together",
               "matcher": hrules.match_connected_hboxes, 
               "rule": hrules.fuse_hboxes, 
               "type": MATCHES_EDGES},
    "mult_hbox": {"text": "Multiply H-boxes", 
               "tooltip": "Merges groups of H-boxes that have the same connectivity",
               "matcher": hrules.match_par_hbox, 
               "rule": hrules.par_hbox, 
               "type": MATCHES_VERTICES},
    "fuse_w": {"text": "fuse W nodes",
               "tooltip": "Merges two connected W nodes together",
               "matcher": rules.match_w_fusion_parallel,
                "rule": rules.w_fusion,
                "type": MATCHES_EDGES},
    "copy": {"text": "copy 0/pi spider", 
               "tooltip": "Copies a single-legged spider with a 0/pi phase through its neighbor",
               "matcher": hrules.match_copy, 
               "rule": hrules.apply_copy, 
               "type": MATCHES_VERTICES},
    "pauli": {"text": "push Pauli", 
               "tooltip": "Pushes an arity 2 pi-phase through a selected neighbor",
               "matcher": pauli_matcher, 
               "rule": pauli_push, 
               "type": MATCHES_VERTICES},
    "bialgebra": {"text": "bialgebra",
               "tooltip": "Applies the bialgebra rule to a connected pair of Z and X spiders",
               "matcher": match_bialgebra,
               "rule": bialgebra,
               "type": MATCHES_EDGES},
    "bialgebra_op": {"text": "bialgebra (inverse)",
               "tooltip": "Applies the bialgebra rule to a connected pair of Z and X spiders in the opposite direction",
               "matcher": match_bialgebra_op,
               "rule": bialgebra_op,
               "type": MATCHES_VERTICES},
    "euler": {"text": "decompose Hadamard", 
               "tooltip": "Expands a Hadamard-edge into its component spiders using its Euler decomposition",
               "matcher": match_hadamard_edge, 
               "rule": euler_expansion, 
               "type": MATCHES_EDGES},
    "lcomp": {"text": "local complementation", 
               "tooltip": "Deletes a spider with a pi/2 phase by performing a local complementation on its neighbors",
               "matcher": rules.match_lcomp_parallel, 
               "rule": rules.lcomp, 
               "type": MATCHES_VERTICES},
    "pivot": {"text": "pivot", 
               "tooltip": "Deletes a pair of spiders with 0/pi phases by performing a pivot",
               "matcher": lambda g, matchf: rules.match_pivot_parallel(g, matchf, check_edge_types=True), 
               "rule": rules.pivot, 
               "type": MATCHES_EDGES}
}


def operations_to_js() -> str:
    global operations
    return json.dumps({k:
            {"active":False, 
             "text":v["text"], 
             "tooltip":v["tooltip"]
            } for k,v in operations.items()})
