from __future__ import nested_scopes
import sys, random
from math import *
from numarray import *
import numarray.linear_algebra as la


class QuValue:
    def __init__(self, repr):
        self.repr = repr
    def __repr__(self):
        return self.repr
    def __cmp__(self, other):
        return cmp(self.repr, `other`)
    def __hash__(self):
        return hash(self.repr)


F = QuValue('F')   # spin up
T = QuValue('T')   # spin down
F.neg = T
T.neg = F


class QuReg:
    counter = 0
    
    def __init__(self, state=F):
        if type(state) == type({}):
            state = state.items()
            values = [value for value, prob in state]
            probs = array([prob for value, prob in state])
            matrix = multiply.outer(probs, probs)
            self.system = Density([self], [values], matrix)
        else:
            self.system = Density([self], [[state]])
        self.name = 'QuReg%d' % QuReg.counter
        QuReg.counter += 1

    def combine(self, other):
        if self.system is not other.system:
            new = self.system.tensor(other.system)
            for reg in new.coords:
                reg.system = new

    def measure(self, measureop=None, initialvalue=F):
        measureop = measureop or fanout
        r_measure = QuReg(initialvalue)
        measureop(self, r_measure)
        s = self.system.trace_keep(r_measure)
        return s.measure()[0]

    def force(self, value):
        system = self.system
        coords = list(system.coords)
        indices = list(system.indices)
        k = coords.index(self)
        oldi = indices[k]
        del coords[k]
        del indices[k]

        ix = oldi.index(value)
        before = prodlen(indices[:k])
        after = prodlen(indices[k:])
        matrix = system.getmatrix(before, len(oldi), after)
        matrix = matrix[:,ix,:,:,ix,:]
        new = Density(coords, indices, matrix.copy())
        new.renormalize()
        for reg in coords:
            reg.system = new
        self.system = Density([self], [[oldi[ix]]])

    def trace_away(self):
        return self.system.trace_away(self)

    def trace_keep(self):
        return self.system.trace_keep(self)

    def forget(self):
        new = self.system.trace_away(self)
        for reg in self.system.coords:
            reg.system = new
        self.system = Forgotten()

    def addvalue(self, value):
        system = self.system
        k = system.coords.index(self)
        if value not in system.indices[k]:
            matrix = system.getmatrixnd()
            base = len(system.indices)
            nlst = list(system.indices)
            rowlst = list(nlst)
            rowlst[k] = [value]
            nlst[k] = list(nlst[k]) + [value]
            l = [len(i) for i in rowlst + system.indices]
            matrix = concatenate((matrix, zeros(l)), k)
            l = [len(i) for i in nlst + rowlst]
            matrix = concatenate((matrix, zeros(l)), base+k)
            new = Density(system.coords, nlst, matrix.copy())
            for reg in system.coords:
                reg.system = new

    def __repr__(self):
        return self.name

    def __str__(self):
        if isinstance(self.system, Forgotten):
            return 'forgotten ' + self.name
        else:
            return self.name + ' in ' + self.system.__repr__(self)

def displayhook(o, hook=sys.displayhook):
    if isinstance(o, QuReg):
        import __builtin__
        print str(o)
        __builtin__._ = o
    else:
        hook(o)
sys.displayhook = displayhook


class Forgotten:
    pass


def prodlen(indices):
    return multiply.reduce([len(i) for i in indices])

def cartesian_product(lst):
    result = []
    def enum(previous, next):
        if not next:
            result.append(previous)
        else:
            next1 = next[1:]
            for i in next[0]:
                enum(previous + (i,), next1)
    enum((), lst)
    return result


class Density:
    
    def __init__(self, coords, indices, matrix=None):
        self.coords = coords
        self.indices = indices
        self.dim = prodlen(indices)
        if matrix is None:
            matrix = zeros((self.dim, self.dim), type = Complex64)
            matrix[0, 0] = 1
        self.matrix = matrix

    def getmatrix(self, *shape):
        matrix = self.matrix
        matrix.setshape(shape*2)
        return matrix

    def getmatrix2d(self):
        return self.getmatrix(self.dim)

    def getmatrixnd(self):
        return self.getmatrix(*[len(i) for i in self.indices])

    def renormalize(self):
        tr = trace(self.getmatrix2d())
        self.matrix /= tr.real

    def tensor(self, other):
        a1 = self.getmatrix2d()[:,NewAxis,:,NewAxis]
        a2 = other.getmatrix2d()[NewAxis,:,NewAxis,:]
        return Density(self.coords + other.coords,
                       self.indices + other.indices, a1 * a2)

    def trace_away(self, *rs):
        while rs:
            r1 = rs[0]
            coords = list(self.coords)
            indices = list(self.indices)
            k = coords.index(r1)
            oldi = indices[k]
            del coords[k]
            del indices[k]
            before = prodlen(indices[:k])
            after = prodlen(indices[k:])
            matrix = self.getmatrix(before, len(oldi), after)
            blockdiagonal = [matrix[:,i,:,:,i,:] for i in range(len(oldi))]
            trace = add.reduce(blockdiagonal)
            self = Density(coords, indices, trace)
            rs = rs[1:]
        return self

    def trace_keep(self, r1):
        coords = self.coords
        indices = self.indices
        k = coords.index(r1)
        dim = len(indices[k])
        before = prodlen(indices[:k])
        beforeindices = range(before)
        after = prodlen(indices[k+1:])
        afterindices = range(after)
        matrix = self.getmatrix(before, dim, after)
        trace = zeros((dim, dim), type = Complex64)
        for ibefore in beforeindices:
            for iafter in afterindices:
                trace += matrix[ibefore,:,iafter,ibefore,:,iafter]
        return Density([r1], [indices[k]], trace)

    def is_classical(self):
        matrix = self.getmatrix2d()
        return allclose(identity(self.dim) * matrix, matrix)

    def force(self, i0):
        assert len(i0) == len(self.coords)
        for value, r1 in zip(i0, self.coords):
            r1.force(value)

    def getindices2d(self):
        return cartesian_product(self.indices)

    def measure(self):
        assert self.is_classical(), "density matrix is not classical"
        m = random.random()
        trace = 0.0
        for indices in self.getindices2d():
            if m >= trace:
                result = indices
            trace += self[indices, indices].real
        assert abs(trace - 1) < 1E-7, (
            "density matrix trace should be 1, it is %s" % trace)
        print 'Measuring', result, 'in', self
        self.force(result)
        return result

    def __getitem__(self, (i, j)):
        matrix = self.getmatrixnd()
        p = []
        q = []
        try:
            for idx, i1, j1 in zip(self.indices, i, j):
                p.append(idx.index(i1))
                q.append(idx.index(j1))
        except ValueError:
            return 0
        return matrix[tuple(p+q)]

##    def __repr__(self, emphasis=None):
##        matrix = self.getmatrix2d()
        
##        lines = str(matrix).split('\n')
##        result = []
##        for line, index in zip(lines, reprindices):
##            result.append('%12s %s' % (index, line))
##        return '\n'.join(result)
        
##        indices = self.getindices2d()
##        assert len(lines) == self.dim == len(indices)
##        reprindices = []
##        for index in indices:
##            repr1 = []
##            for i, r in zip(index, self.coords):
##                s = `i`
##                if r is not emphasis:
##                    s = s.lower()
##                repr1.append(s)
##            reprindices.append(','.join(repr1))

##        eval, evec = la.Heigenvectors(matrix)
        
##        ...

        
##        result = []
##        for line, index in zip(lines, reprindices):
##            result.append('%12s %s' % (index, line))
##        return '\n'.join(result)

    def __repr__(self, emphasis=None):
        matrix = self.getmatrix2d()
        indices = self.getindices2d()
        val2i = dict(zip(indices, range(len(indices))))
        indices.sort()
        reprindices = []
        for i in indices:
            repr1 = []
            for v, r in zip(i, self.coords):
                s = `v`
                if r != emphasis:
                    s = s.lower()
                repr1.append(s)
            reprindices.append(','.join(repr1))

        def repr1(v):
            if abs(v) < 1E-8:
                return ''
            elif type(v) != type(1j) or abs(v.imag) < 1E-8:
                return '%.3g' % getattr(v, 'real', v)
            else:
                return '%.3g+%.3gj' % (v.real, v.imag)
        
        lines = [[] for i in range(self.dim+2)]
        eval, evec = la.eigenvectors(matrix)  # NB. Heigenvectors() broken??
        for i in range(self.dim):
            val = eval[i]
            vec = evec[i]
            val = getattr(val, 'real', val)
            if abs(val) > 1E-7:
                lines[0].append('')
                if 0.995 <= val <= 1.005:
                    lines[1].append('100%')
                else:
                    lines[1].append('%.2f%%' % (100.0*val))
                for i in range(self.dim):
                    idx = val2i[indices[i]]
                    lines[i+2].append(repr1(vec[idx]))
        if len(lines[0]) == 1:
            header = 'Pure state'
        else:
            header = 'Mixed state'

        lines[0].append('')
        lines[0].append('')
        lines[0].append('|')
        for i in range(self.dim):
            lines[0].append(reprindices[i])
            idx1 = val2i[indices[i]]
            line = lines[i+2]
            line.append(' |')
            line.append(reprindices[i])
            line.append('|')
            for j in indices:
                line.append(repr1(matrix[idx1, val2i[j]]))
        
        columns = {}
        for line in lines:
            for i in range(len(line)):
                if len(line[i]) >= columns.get(i, 0):
                    columns[i] = len(line[i])
        coords = [r.name for r in self.coords]
        strlines = ['%s [%s]' % (header, ','.join(coords))]
        for line in lines:
            for i in range(len(line)):
                line[i] = line[i].center(columns[i])
            strlines.append(' '.join(line))
        strlines.append('')
        pos = strlines[3].index('|')
        strlines[2] = strlines[2].ljust(pos) + ',' + '-'*(len(strlines[3])-pos-2)
        pos = strlines[3].rfind('|')
        strlines[2] = strlines[2][:pos] + '+' + strlines[2][pos:]
        return '\n'.join(strlines)


##def extendmatrix(matrix, srcdims, targetlens):
##    shape = [targetlens[d] for d in srcdims]
##    axemap = list(srcdims)
##    for i in range(len(targetlens)):
##        if i not in srcdims:
##            shape.insert(i, 1)
##            axemap.insert(i, None)
##    matrix.setshape(shape)
##    axes = range(len(shape))
##    for i in range(len(shape)):
##        if axemap[i] is not None:
##            axes[i] = axemap[i]
##    matrix = transpose(matrix, axes)
##    for i in range(len(shape)):
##        if axemap[i] is None:
##            matrix = repeat(matrix, array([targetlens[i]]), i)
##    return matrix

def extendmatrix(matrix, srcdims, targetlens):
    shape1 = [targetlens[d] for d in srcdims]
    shape2 = []
    axemap = list(srcdims)
    missing_sqrt = 1
    for i in range(len(targetlens)):
        if i not in srcdims:
            shape1.insert(i, 1)
            axemap.insert(i, None)
            if 2*i < len(targetlens):
                missing_sqrt *= targetlens[i]
            shape2.append(targetlens[i])
        else:
            shape2.append(1)
    matrix.setshape(shape1)
    axes = range(len(shape1))
    for i in range(len(shape1)):
        if axemap[i] is not None:
            axes[i] = axemap[i]
    matrix = transpose(matrix, axes)
    id = identity(missing_sqrt)
    id.setshape(shape2)
    return matrix * id


class Operation:
    _cache_key = None

    def __init__(self, table=None):
        self.table = table

    def f(self, *vs):
        if len(vs) == 1:
            vs, = vs
        return self.table[vs]

    def unitarymatrix(self, inputidx):
        inputindices = cartesian_product(inputidx)
        outputidx = inputidx
        outputindices = inputindices
        while 1:
            dim = len(outputindices)
            val2i = dict(zip(outputindices, range(dim)))
            matrix = zeros((dim, len(inputindices)), type = Complex64)
            try:
                for col in range(len(inputindices)):
                    invalue = inputindices[col]
                    for seq in self.f(*invalue):
                        assert len(seq) == len(inputidx)+1
                        value = tuple(seq[:-1])
                        row = val2i[value]   # possibly KeyError
                        matrix[row,col] += seq[-1]
                return outputidx, matrix
            except KeyError:
                noutputidx = []
                for newi, iold in zip(value, outputidx):
                    if newi not in iold:
                        iold = iold + [newi]
                    noutputidx.append(iold)
                outputidx = noutputidx
            outputindices = cartesian_product(outputidx)

    def extendedmatrix(self, system, ks):
        key = system.indices, ks
        if key != self._cache_key:
            ks2 = [k + len(system.indices) for k in ks]
            inputidx = [system.indices[k] for k in ks]
            outputidx, U = self.unitarymatrix(inputidx)
            # extend U to cover the extra dimensions of A
            rows = list(system.indices)
            for k, inew in zip(ks, outputidx):
                rows[k] = inew
            newdim = prodlen(rows)
            U = extendmatrix(U, ks+ks2, [len(row) for row in rows] +
                                        [len(col) for col in system.indices])
            U = reshape(U, (newdim, system.dim))
            self._cache_key = key
            self._cache_result = U, rows, newdim
        return self._cache_result

    def __call__(self, *rs):
        r1 = rs[0]
        for r in rs[1:]:
            r1.combine(r)
        system = r1.system
        ks = []
        for r in rs:
            k = system.coords.index(r)
            assert k not in ks, (
                "%s on duplicated QuReg arguments" % self.__class__.__name__)
            ks.append(k)
        U, rows, newdim = self.extendedmatrix(system, ks)
        A = system.getmatrix2d()
        Ustar = conjugate(transpose(U))
        # U * A * Ustar
        B = matrixmultiply(U, matrixmultiply(A, Ustar))
        system.indices = rows
        system.dim = newdim
        system.matrix = B


def rotate(r1, angle):
    sina, cosa = sin(angle), cos(angle)
    op = Operation({
        F: [(F,  cos(angle)), (T,  sin(angle))],
        T: [(F, -sin(angle)), (T,  cos(angle))],
        })
    return op(r1)

neg = Operation({
    F: [(T,1)],
    T: [(F,1)],
    })

fanout = Operation({
    (F,F): [(F,F,1)],
    (F,T): [(F,T,1)],
    (T,F): [(T,T,1)],
    (T,T): [(T,F,1)],
    })


try:
    f = open('/dev/urandom', 'rb')
    random.seed(f.read(16))
    f.close()
except IOError:
    pass


r0 = QuReg()
#print r0
#for i in range(4):
#    apply1(rotate, r0)
#    print r0

rotate(r0, pi/4)
r1 = QuReg()
r2 = QuReg()
r0.combine(r1)
r1.addvalue(T)
r0.combine(r2)
r2.addvalue(T)
#print r1
fanout(r0, r1)
fanout(r1, r2)
#print r2

#print r0.measure()
#print r0
#STOP

#print r0
##print 'Measure:', r0.measure()
##print 'Measure:', r1.measure()
##print 'Measure:', r2.measure()
##print 'Measure:', r3.measure()
##print 'Measure:', r4.measure()
##print 'Measure:', r5.measure()
#print r5
#print r0


class Succ(Operation):
    
    def f(self, v):
        return [(v+1, 1)]
        
succ = Succ()


class IRotate(Operation):
    
    def __init__(self, angle, n1, n2):
        self.n1 = n1
        self.n2 = n2
        self.sina, self.cosa = sin(angle), cos(angle)
        
    def f(self, v):
        if v == self.n1:
            return [(self.n1, self.cosa),
                    (self.n2, -self.sina)]
        elif v == self.n2:
            return [(self.n1, self.sina),
                    (self.n2, self.cosa)]
        else:
            return [(v, 1)]


class AtLeast(Operation):

    def __init__(self, n):
        self.n = n

    def f(self, v, w):
        if v < self.n:
            return [(v, w, 1)]
        else:
            return [(v, w.neg, 1)]


class FanoutInt(Operation):

    def f(self, v, w):
        return [(v, v^w, 1)]

fanoutint = FanoutInt()


##s0 = QuReg(0)
##for i in range(5):
##    #print s0
##    op = IRotate(pi/4, i, i+1)
##    op(s0)
##print s0
##s1 = QuReg(0)
##fanoutint(s0, s1)
##s0.measure(AtLeast(1))
##s1.measure(AtLeast(2))
##print s0
##print s1
##s0.measure(fanoutint, 0)
##STOP

class CondRotate(Operation):
    def __init__(self, n):
        self.n = n
    def f(self, n):
        if n == self.n:
            return [(n, -1)]
        else:
            return [(n, 1)]

N = 5

diffusion_table = {}
for n in range(N):
    diffusion_table[n] = [(n, -1)] + [(j, 2.0/N) for j in range(N)]


d = {}
for n in range(N):
    d[n] = 1.0/sqrt(N)
s0 = QuReg(d)
#print s0

condrotate = CondRotate(2)
diffusion = Operation(diffusion_table)

M = int(pi/4*sqrt(N))
for i in range(M):
    print 'step', i, '/', M, '...'
    condrotate(s0)
    #print s0
    diffusion(s0)
    #print s0
print s0.measure(fanoutint, 0)


# base-2 equivalent

TARGET = (F, T, T, F, T, T)
N = len(TARGET)

class CondRotate(Operation):
    def f(self, *n123):
        if n123 == TARGET:
            return [n123 + (-1,)]
        else:
            return [n123 + (1,)]

def cartesian_power(set, n):
    if n == 0:
        return [()]
    else:
        c = cartesian_power(set, n-1)
        return [head + (a,) for head in c for a in set]

states = {}
for n123 in cartesian_power([F, T], N):
    states[n123] = len(states)

diffusion_table = {}
tail = []
for n123 in states:
    tail.append(n123 + (2.0/(2**N),))
for n123 in states:
    diffusion_table[n123] = [n123 + (-1,)] + tail

hadamard_table = {}
f = 2 ** (-0.5*N)
def parity(n):
    if n&1:
        return -1
    else:
        return 1
for state, num in states.items():
    hadamard_table[state] = [state2+(parity(num&num2)*f,)
                             for state2, num2 in states.items()]

s = [QuReg() for i in range(N)]

hadamard = Operation(hadamard_table)
condrotate = CondRotate()
diffusion = Operation(diffusion_table)

hadamard(*s)
#print s[0]

if 1:
    M = int(pi/4*sqrt(2**N))
    for i in range(M):
        print 'step', i, '/', M, '...'
        condrotate(*s)
        #print s[0]
        diffusion(*s)
        #print s[0]
    m = [s0.measure() for s0 in s]
    print m