from numpy import *
from openopt import SLE
import scipy.sparse.linalg

class Sketch(object):
    def __init__(self):
        self.shapes = []
        self.constraints = []

    def __iter__(self):
        return self.shapes.__iter__()

    def __contains__(self, item):
        return self.shapes.__contains__(item)

    def solve(self):
        shape_vars = {}
        varid_to_shape_var = []
        varid = 0
        equations = []
        for c in self.constraints:
            equations += c.toEqu()
        for e in equations:
            coeffs = e[0]
            for c in coeffs:
                shape_var = c[1]
                if not shape_var in shape_vars:
                    shape_vars[shape_var] = varid
                    varid += 1
                    varid_to_shape_var.append(shape_var)
        dim = max(len(equations), len(shape_vars))
        C = zeros((dim, dim))
        d = zeros(dim)
        for i in range(len(equations)):
            e = equations[i]
            coeffs = e[0]
            for c in coeffs:
                var_index = shape_vars[c[1]]
                C[i][var_index] = c[0]
            d[i] = e[1]
#        for row in range(len(equations)):
#            print '[',
#            for col in range(len(shape_vars)):
#                print C[row][col],
#                print ',',
#            print '] [%f]' % d[row]
        s = SLE(C, d)
        r = s.solve()
        solved_vars = r.xf
        print "SOLVED: ", solved_vars
        for i in range(len(solved_vars)):
            varid_to_shape_var[i][0].setVar(varid_to_shape_var[i][1],
                    solved_vars[i])