#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Apr  6 11:53:08 2020

@author: pemeriau
"""

from pyomo.environ import *
import numpy as np

##############################
# Test (2,1) RAC.
# 2 input bits encoded in 1 qubit
# Find best {states,POVM} that maximises the retrieval of one of the bit
# Expected output : .85
##############################

### Inputs
dim = 2 #dimension the problem
nb_input = 2 #number of inputs
size_input = 2 #size of the inputs (2=bits,3=trits...)
nb_comp = 4 #number of components: 4 because hermitian
            


### create the  model
model = ConcreteModel()

### Initialise sets
model.size = RangeSet(1,nb_comp)
model.inputs = RangeSet(0,nb_input-1)
model.x = RangeSet(0,size_input-1)
# Size taking into account number of inputs and their size
for i in range(nb_input):
    model.size *= model.x

# model.size of the form:
#   [n,x0,x1] : n in {1,2,3,4} and x0,x1 in {0,1} for bits
#   |  |  |
#   |  ---------------x0---------x1
# which component     |          |
# a b c or d in :  value of x0   |
# [a c+id                    value of x1
#  c-id b]
# as we are dealing with hermitian matrices


### Variable for p (the probability to maximise)
#model.p = Var()
model.p = Var(bounds=(0.0,1.0))

### Variables for density matrices (real & imaginary)
model.rho = Var(model.size)

### Variables for POVM (real & imginary)
model.E = Var(model.size)


### Define objective
model.OBJ = Objective(expr = model.p, sense=maximize)

### 1st constraint : Tr(E*rho) >= p forall x and i
#one for each x in {0,1}^2 and for each i
#for qubits
def bilinear_const(model,x0,x1,i):
    if i == 0:
        trace = model.rho[1,x0,x1]*model.E[1,i,x0] \
            + 2* model.rho[3,x0,x1]*model.E[3,i,x0] \
            + 2* model.rho[4,x0,x1]*model.E[4,i,x0] \
            + model.rho[2,x0,x1]*model.E[2,i,x0]
    if i == 1:
        trace = model.rho[1,x0,x1]*model.E[1,i,x1] \
            + 2* model.rho[3,x0,x1]*model.E[3,i,x1] \
            + 2* model.rho[4,x0,x1]*model.E[4,i,x1] \
            + model.rho[2,x0,x1]*model.E[2,i,x1]
    return model.p <= trace

model.bilinear_constraint = Constraint(model.x,model.x,model.inputs,rule=bilinear_const)

### 2nd constraint: Tr(rho) == 1 forall x
def unit_trace(model,x0,x1):
    trace = model.rho[1,x0,x1]+model.rho[2,x0,x1]
    return trace == 1

model.unit_trace_constraint = Constraint(model.x,model.x,rule=unit_trace)

### 3rd constraint: E_0 + E_1 == Id forall i
def POVM_1(model,i):
    return (model.E[1,i,0] + model.E[1,i,1]) == 1
def POVM_2(model,i):
    return (model.E[2,i,0] + model.E[2,i,1]) == 1
def POVM_3(model,i):
    return (model.E[3,i,0] + model.E[3,i,1]) == 0
def POVM_4(model,i):
    return (model.E[4,i,0] + model.E[4,i,1]) == 0

model.POVM1 = Constraint(model.inputs,rule=POVM_1)
model.POVM2 = Constraint(model.inputs,rule=POVM_2)
model.POVM3 = Constraint(model.inputs,rule=POVM_3)
model.POVM4 = Constraint(model.inputs,rule=POVM_4)

### 4th constraint: SDP
def SDP_rho1(model,x0,x1):
    return model.rho[1,x0,x1] >= 0
def SDP_rho2(model,x0,x1):
    discr = model.rho[1,x0,x1]*model.rho[2,x0,x1] \
        - model.rho[3,x0,x1]**2 - model.rho[4,x0,x1]**2
    return discr  >= 0
def SDP_E1(model,i,x):
    return model.E[1,i,x] >= 0
def SDP_E2(model,i,x):
    discr = model.E[1,i,x]*model.E[2,i,x] \
        - model.E[3,i,x]**2 - model.E[4,i,x]**2
    return discr  >= 0

model.SDPrho1 = Constraint(model.x,model.x,rule=SDP_rho1)
model.SDPrho2 = Constraint(model.x,model.x,rule=SDP_rho2)
model.SDPE1 = Constraint(model.inputs,model.x,rule=SDP_E1)
model.SDPE2 = Constraint(model.inputs,model.x,rule=SDP_E2)