import numpy as np
import picos, math
from itertools import product
import numpy as np
import picos, math
from itertools import product
from igraph import *

yo13 = [[1,0,0],[0,1,0],[0,0,1],[0,1,-1],[1,0,-1],[1,-1,0],[0,1,1],[1,0,1],[1,1,0],[-1,1,1],[1,-1,1],[1,1,-1],[1,1,1]]
Id = np.eye(3)
pmos = [Id-2*np.outer(yo13[i],yo13[i])/np.dot(yo13[i],yo13[i]) for i in range(13)] # ideal single operators

def assignment(case,contexts):
    res = []
    for i, context in enumerate(contexts):
        tmp = [case[j] for j in context]
        res.append(math.prod(tmp))
    return res

def edobservable(A,B,C,eta,sa,sb,sc):
    Id = np.eye(3)
    term1 = eta**3 * np.dot(np.dot(A, B), C)
    term2 = eta**2 * (1 - eta) * np.dot(A, B) * sc
    term3 = eta * (1 - eta) * A * sb * sc
    term4 = (1 - eta) * Id * sa * sb * sc
    return term1 + term2 + term3 + term4
# we have to modify this part such that A -> epsilon A instead of epsilon A + (1-epsilon) I
# eta A + delta_A I, where delta_A can be 1 or -1

def edobservable2(As,eta,ss): #complete the list such that the length is 3
    Id = np.eye(3)
    Bs= As + [Id]*3
    ts = ss + [1]*3
    A,B,C = Bs[:3]
    sa,sb,sc = ts[:3] 
    return edobservable(A,B,C,eta,sa,sb,sc)
    
def violation(eta,epsilon,spms):
    orderedOps = [edobservable2([(1-2*epsilon)*pmos[i] for i in context], eta, [spms[i] for i in context]) for context in contexts]
    A = np.matrix([np.array(op).flatten().tolist() for op in orderedOps])
    b = np.matrix([1,0,0,0,1,0,0,0,1]).T
    x = picos.RealVariable("x", nConts)
    z = picos.RealVariable("z", 1)
    P = picos.Problem()
    P.set_objective('min', z)
    P.add_list_of_constraints(picos.sum(vectors[j][i]*x[i] for i in range(nConts))<=z for j in range(2**13))
    P.add_list_of_constraints(picos.sum(A[i,j]*x[i] for i in range(nConts))==b[j] for j in range(9))
    sol = P.solve(solver="mosek")
    V = (1/sol.value.real-1)
    return V

#print(violation(0.99,0,[-1]*10+[1]*3))
def vec2proj(v,r=1):
    if r==1:
        return np.outer(v,v)/np.dot(v,v)
    else:
        proj = np.outer(v,np.conj(v))/np.dot(v,np.conj(v))
        projr = np.block([[proj.real,proj.imag],[-1*proj.imag,proj.real]])
        return projr

def vec2eds(sicset,r=1):
    projs = [vec2proj(v,r) for v in sicset]
    d = len(sicset[0])
    if r != 1: d = 2*d
    identity = np.identity(d)
    size = len(sicset)
    eds = [[i,j] for i in range(size) for j in range(i,size) if abs(np.trace(np.dot(projs[i],projs[j])))<=1e-5]
    #eds = [[i,j] for i in range(size) for j in range(i,size) if np.dot(sicset[i],sicset[j])==0]
    return eds

def vec2gh(sicset,r=1):
    eds = vec2eds(sicset,r)
    gh = Graph(eds)
    return gh

# run over all possible eta and epsilon to find out their relations for a fixed ineq
def violation2(eta,epsilon,contexts,ineq,vectors,spms):
    orderedOps = [edobservable2([(1-2*epsilon)*pmos[i] for i in context], eta, [spms[i] for i in context]) for context in contexts]
    A = np.matrix([np.array(op).flatten().tolist() for op in orderedOps])
    res = sum([coe*orderedOps[i] for i,coe in enumerate(ineq)])
    classical = max([np.dot(it,ineq) for it in vectors])
    return [np.linalg.eigvals(res),classical]
    
def get_canonical_form(lst, group_perms):
    arr = np.array(lst)
    candidates = [arr[p] for p in group_perms]
    return tuple(min(candidates, key=lambda x: list(x)))

def remove_automorphic_duplicates(sets_of_lists, group_perms):
    unique_canonical_forms = set()
    result = []
    for lst in sets_of_lists:
        canonical = get_canonical_form(lst, group_perms)
        if canonical not in unique_canonical_forms:
            unique_canonical_forms.add(canonical)
            result.append(lst)        
    return result

def ob4cont(subc,cont,eta,epsilon,spms):
    As = [(1-2*epsilon)*pmos[t] for t in cont]
    ss = [spms[t] for t in cont]
    Id = np.eye(3)
    pos = [cont.index(i) for i in set(cont)-set(subc)]
    for t in pos:
        As[t] = Id
        ss[t] = 1
    return edobservable2(As,eta,ss)

# figure 3 in the paper: for given eta, find critical epsilon
def curve(eta,epsilon0,ineq,contexts,spms):
    first_cols = list(product([-1, 1], repeat=13))
    vectors = [assignment(case,contexts) for case in first_cols]
    epsilonu = epsilon0
    epsilonl = 0
    epsilon = (epsilonu+epsilonl)/2
    qs,cb = violation2(eta,0,contexts,ineq,vectors,spms)
    if min(qs)-cb < 0: return -1
    qs,cb = violation2(eta,epsilon,contexts,ineq,vectors,spms)
    while abs(min(qs)-cb) > 1e-4:
        if min(qs)-cb >= 0: 
            epsilonl = epsilon
            epsilon = (epsilonl+epsilonu)/2
        else: 
            epsilonu = epsilon
            epsilon = (epsilonl+epsilonu)/2
        qs,cb = violation2(eta,epsilon,contexts,ineq,vectors,spms)
    return epsilon

# figure 2 in the paper: for a given eta, find the critical epsilon with the consideration of disturbance
def curve2(eta,epsilon0,eevents3,cb,cd,spms): # cb is classical bound, cd is coefficient of disturbance
    epsilonu = epsilon0
    epsilonl = 0
    epsilon = (epsilonu+epsilonl)/2
    q0 = min(np.linalg.eigvals(sum([coe*ob4cont(subc,cont,eta,epsilon,spms) for coe,subc,cont in eevents3])))
    delta2 = max(abs(np.linalg.eigvals(ob4cont([5],[5,2,8],eta,epsilon,spms)-ob4cont([5],[11,5],eta,epsilon,spms))))+max(abs(np.linalg.eigvals(ob4cont([4],[1,7,4],eta,epsilon,spms)-ob4cont([4],[10,4],eta,epsilon,spms))))+max(abs(np.linalg.eigvals(ob4cont([6],[6,9,3],eta,epsilon,spms)-ob4cont([6],[12,6],eta,epsilon,spms))))+max(abs(np.linalg.eigvals(ob4cont([8],[5,2,8],eta,epsilon,spms)-ob4cont([8],[10,8],eta,epsilon,spms))))
    while abs(q0-cd*delta2-cb) > 1e-4:
        if q0-cd*delta2-cb >= 0: 
            epsilonl = epsilon
            epsilon = (epsilonl+epsilonu)/2
        else: 
            epsilonu = epsilon
            epsilon = (epsilonl+epsilonu)/2
        #print(epsilon)
        q0 = min(np.linalg.eigvals(sum([coe*ob4cont(subc,cont,eta,epsilon,spms) for coe,subc,cont in eevents3])))
        delta2 = max(abs(np.linalg.eigvals(ob4cont([5],[5,2,8],eta,epsilon,spms)-ob4cont([5],[11,5],eta,epsilon,spms))))+max(abs(np.linalg.eigvals(ob4cont([4],[1,7,4],eta,epsilon,spms)-ob4cont([4],[10,4],eta,epsilon,spms))))+max(abs(np.linalg.eigvals(ob4cont([6],[6,9,3],eta,epsilon,spms)-ob4cont([6],[12,6],eta,epsilon,spms))))+max(abs(np.linalg.eigvals(ob4cont([8],[5,2,8],eta,epsilon,spms)-ob4cont([8],[10,8],eta,epsilon,spms))))
    return [eta,epsilon]

# symmetry in YO13
gh = vec2gh(yo13)
automorphisms = gh.get_automorphisms_vf2()
signs = list(product([-1, 1], repeat=13))
ssigns = remove_automorphic_duplicates(signs, automorphisms)
len(ssigns)