from numpy import *
from openopt import SLE

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)
        C = empty((len(equations), len(shape_vars)))
        d = empty(len(equations))
        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]
        s = SLE(C, d)
        r = s.solve()
        solved_vars = r.xf
        for i in range(len(solved_vars)):
            varid_to_shape_var[i][0].setVar(varid_to_shape_var[i][1],
                    solved_vars[i])