# Copyright (C) 2023 Daniel Gerbet

#from sage.all import PolynomialRing
#from sage.all import AA
from sage.all import *

class Rect:
    def __init__(self, row, col, rows, cols):
        self.row = row
        self.col = col
        self.rows = rows
        self.cols = cols

    def __repr__(self):
        return "rect(%d,%d)++(%d,%d)"%(self.row, self.col, self.rows, self.cols)

# TODO
# build all matrices with increasing dimension
# while ensuring rows >= cols:

# a rect has 4 corners
# 4 corners also has the square
# for each split, horizontal or vertical,
# there are at least 4 corners of the rects adjacent

class RatioSet:
    def __init__(self, item=None):
        if item is None:
            self.set = []
        else:
            self.set = [item]

    def __repr__(self):
        return "(%d) %s"%(len(self.set), self.set)

    def __len__(self):
        return len(self.set)

    def __bool__(self):
        return len(self.set) > 0

    def add(self, item):
        # TODO: use bijection
        i = 0
        while i < len(self.set) and self.set[i] < item:
            i += 1
        if i < len(self.set):
            if self.set[i] == item:
                return
            self.set = self.set[:i] + [item] + self.set[i:]
        else:
            self.set.append(item)

    def union(self, other):
        ans = RatioSet()
        i = 0
        j = 0
        while i < len(self.set) and j < len(other.set):
            if self.set[i] < other.set[j]:
                ans.set.append(self.set[i])
                i += 1
                continue
            if other.set[j] < self.set[i]:
                ans.set.append(other.set[j])
                j += 1
                continue
            ans.set.append(self.set[i])
            i += 1
            j += 1
        ans.set += self.set[i:]
        ans.set += other.set[j:]
        return ans

class RatioDict:
    def __init__(self, n=None, s=None):
        if n is None:
            self.dt = []
        else:
            assert n >= 1
            self.dt = [RatioSet()]*(n-1) + [s]

    def __repr__(self):
        if not self.dt:
            return "empty"
        ans = ""
        for n, s in enumerate(self.dt):
            if s:
                ans += "%d : %s\n"%(n+1, s)
        return ans

    def __getitem__(self, idx):
        if idx-1 >= len(self.dt):
            return RatioSet()
        return self.dt[idx-1]

    def union(self, other):
        ans = RatioDict()
        ans.dt = [s.union(o) for s, o in zip(self.dt, other.dt)]
        l = len(ans.dt)
        ans.dt += self.dt[l:]
        ans.dt += other.dt[l:]
        return ans


def SquareDivision(rows, cols, max_nrect=None):
    sd = _SquareDivisionNode(None, None, (rows, cols))
    sd._complete(max_nrect=max_nrect)
    return sd

class _SquareDivisionNode:
    def __init__(self, parent, rect, dim=None):
        self.parent = parent
        self.rect = rect

        if parent is None:
            self.rows = dim[0]
            self.cols = dim[1]
            self.mat = [[0]*(self.cols+1)]*(self.rows+1)
            self.row_split = [0]*self.rows
            self.col_split = [0]*self.cols
            self.nrect = 0
        else:
            self.rows = parent.rows
            self.cols = parent.cols
            self.mat = [[parent.mat[i][j] for j in range(self.cols+1)] for i in range(self.rows+1)]
            self.row_split = [i for i in parent.row_split]
            self.col_split = [i for i in parent.col_split]
            self.nrect = parent.nrect + 1

            # fill the rect
            assert rect.row >= 0
            assert rect.col >= 0
            assert rect.rows > 0
            assert rect.cols > 0
            assert rect.row+rect.rows <= self.rows+1
            assert rect.col+rect.cols <= self.cols+1
            for i in range(rect.row, rect.row+rect.rows):
                for j in range(rect.col, rect.col+rect.cols):
                    assert self.mat[i][j] == 0
                    self.mat[i][j] = self.nrect
            if rect.row > 0:
                assert self.row_split[rect.row-1]
            if rect.col > 0:
                assert self.col_split[rect.col-1]
            r = rect.row+rect.rows-1
            if r < self.rows:
                self.row_split[r] = 1
            c = rect.col+rect.cols-1
            if c < self.cols:
                self.col_split[c] = 1

    def _complete(self, max_nrect=None):
        assert not hasattr(self, 'complete')
        # search for a free position, where a rect can be inserted
        row = 0
        while row <= self.rows:
            col = 0
            while col <= self.cols:
                if self.mat[row][col] == 0:
                    break
                col += 1
            if col <= self.cols:
                break
            row += 1
        if row > self.rows:
            self.complete = True
            return

        # insert all fitting rects there
        self.complete = False
        self.child = []
        if max_nrect is not None and self.nrect >= max_nrect:
            return
        rows = 1
        while row+rows <= self.rows+1:
            cols = 1
            while col+cols <= self.cols+1:
                r = row
                while r < row+rows:
                    if self.mat[r][col+cols-1]:
                        break
                    r += 1
                if r < row+rows:
                    break
                sd = _SquareDivisionNode(self, Rect(row, col, rows, cols))
                sd._complete(max_nrect=max_nrect)
                if sd._is_fully_splitted():
                    self.child.append(sd)
                cols += 1
            rows += 1

    def __repr__(self):
        ans = " "
        for j in range(self.cols):
            if self.col_split[j]:
                ans += " |"
            else:
                ans += "  "
        ans += "\n"
        for i in range(self.rows+1):
            if i < self.rows and self.row_split[i]:
                ans += "_"
            else:
                ans += " "
            for j in range(self.cols+1):
                if self.mat[i][j]:
                    ans += "%d "%self.mat[i][j]
                else:
                    ans += '. '
            ans += "\n"
        return ans

    def _is_fully_splitted(self):
        if self.complete:
            return 0 not in self.row_split+self.col_split
        return bool(self.child)

    def num_rect_bound(self):
        if self.complete:
            return (self.nrect, self.nrect)
        b = [c.num_rect_bound() for c in self.child]
        l = min(b[0] for b in b) 
        u = max(b[1] for b in b) 
        return (l, u)

    def print_division(self):
        if self.complete:
            print(self)
        else:
            for c in self.child:
                c.print_division()

    def _compute_ratio(self, ideal, rvar, cvar, ratio, max_nrect=None):
        """
        insert the equations that define the ratio for the rect in this stage
        then compute recursively
        """
        def select(var, idx):
            assert idx >= 0 and idx <= len(var)+1
            if idx == 0:
                return 0
            if idx == len(var)+1:
                return 1
            return var[idx-1]

        if max_nrect is not None and self.nrect > max_nrect:
            return

        rl = select(rvar, self.rect.row)
        ru = select(rvar, self.rect.row+self.rect.rows)
        cl = select(cvar, self.rect.col)
        cu = select(cvar, self.rect.col+self.rect.cols)

        ideal_vert = [(o+[0], i + (ratio*(ru-rl) - (cu-cl))) for o, i in ideal]
        ideal_horiz = [(o+[1], i + ((ru-rl) - ratio*(cu-cl))) for o, i in ideal]

        #ideal = set(i for i in ideal_vert + ideal_horiz)
        ideal = ideal_vert + ideal_horiz
        ideal = [(o, i) for o, i in ideal if not i.is_one()]
        #ideal = [i for i in ideal_vert + ideal_horiz if not i.is_one()]
        # TODO: remove duplicates

        if not self.complete:
            for c in self.child:
                c._compute_ratio(ideal, rvar, cvar, ratio, max_nrect=max_nrect)
            return

        # the division of a square into 1 rect case:
        if not rvar+cvar:
            self.ratio = RatioSet(1)
            return

        self.ratio = RatioSet()
        V0 = []
        for o, i in ideal:
            V = i.variety()
            for V in V:
                if V[ratio] <= 0 or V[ratio] > 1:
                    continue
                discard = False
                for x in rvar+cvar:
                    if V[x] <= 0 or V[x] >= 1:
                        discard = True
                        break
                for i in range(len(rvar)-1):
                    if V[rvar[i]] >= V[rvar[i+1]]:
                        discard = True
                for i in range(len(cvar)-1):
                    if V[cvar[i]] >= V[cvar[i+1]]:
                        discard = True
                if discard:
                    continue
                V0.append((o, V))
            if False:
                for V in V0:
                    print(V)
                if V0:
                    print(self)
        for V in V0:
            self.ratio.add(V[1][ratio])
        #print(self.ratio)
        self.variety = V0
        if self.variety:
            self.rvar = rvar
            self.cvar = cvar
            self.ratvar = ratio


    def compute_ratio(self, max_nrect=None):
        """
        compute the ratios of the rectangle sides
        return a dictionary with the number of rects as keys
        and a set of ratios as values

        TODO: allow to compute only for a given set of number of rects
        """
        rvar = ['r%d'%i for i in range(self.rows)]
        cvar = ['c%d'%i for i in range(self.cols)]
        R = PolynomialRing(AA, rvar+cvar+['x'])
        rvar = R.gens()[:self.rows]
        cvar = R.gens()[self.rows:self.rows+self.cols]
        ratio = R.gens()[-1]

        assert not self.complete
        for c in self.child:
            c._compute_ratio([([], Ideal(R, 0))], rvar, cvar, ratio, max_nrect=max_nrect)

    def ratios(self, max_nrect=None):
        """
        get all possible ratios
        """
        if self.complete:
            if self.nrect is not None and self.nrect > max_nrect:
                return RatioDict()
            return RatioDict(self.nrect, self.ratio)
        ratio = RatioDict()
        for c in self.child:
            ratio = ratio.union(c.ratios(max_nrect=max_nrect))
        return ratio

    def find(self, nrect, ratio):
        """
        search for a particular division
        """
        if not self.complete:
            for c in self.child:
                ans = c.find(nrect, ratio)
                if ans is not None:
                    return ans
            return None
        if self.nrect != nrect:
            return None
        if ratio not in self.ratio.set:
            return None
        return self

    def to_tikz(self, file, ratio):
        """
        render a divsion with specified ratio
        """
        assert self.complete
        assert ratio in self.ratio.set

        for V in self.variety:
            if V[1][self.ratvar] == ratio:
                VV = V[1]
                break

        def select(variety, var, idx):
            assert idx >= 0 and idx <= len(var)+1
            if idx == 0:
                return float(0)
            if idx == len(var)+1:
                return float(1)
            return float(variety[var[idx-1]])

        p = self
        while p.rect is not None:
            rl = select(VV, self.rvar, p.rect.row)
            ru = select(VV, self.rvar, p.rect.row+p.rect.rows)
            cl = select(VV, self.cvar, p.rect.col)
            cu = select(VV, self.cvar, p.rect.col+p.rect.cols)
            file.write("\\draw(%s,%s)rectangle(%s,%s);\n"%(rl,cl,ru,cu))
            p = p.parent


if __name__ == "__main__":
    max_nrect = 7
    ratios = RatioDict()
    sd = []
    for rows in range(max_nrect):
        for cols in range(min(rows+1, max_nrect-rows)):
            print("%d rows, %d cols"%(rows, cols))

            sd.append(SquareDivision(rows, cols, max_nrect=max_nrect))
            print(sd[-1].num_rect_bound())

            sd[-1].compute_ratio(max_nrect)
            rat = sd[-1].ratios(max_nrect)
            ratios = ratios.union(rat)

    print(ratios)

    for n in range(2, max_nrect+1):
        f = open("rect%d.tex"%n, "w")
        f.write("\\documentclass[a4paper,10pt]{article}\n")
        f.write("\\usepackage[top=1cm,bottom=1cm,left=1cm,right=1cm]{geometry}\n")
        f.write("\\usepackage{tikz}\n")
        f.write("\\pagestyle{empty}\n")
        f.write("\\begin{document}\n")
        end = True
        max_row = 9
        max_col = 6
        row = 0
        col = 0
        for r in ratios[n].set:
            i = 0
            while i < len(sd):
                s = sd[i].find(n, r)
                if s is not None:
                    break
                i += 1
            #print("%d %s"%(n, r))
            #print(s)

            if end:
                f.write("\\begin{tikzpicture}\n")
                end = False
            f.write("\\begin{scope}[xshift=%scm,yshift=%scm,scale=2]\n"%(col*3, -row*3))
            s.to_tikz(f, r)
            f.write("\\path(.5,0)node[below]{$%.8g$};\n"%float(r))
            f.write("\\end{scope}\n");
            col += 1
            if col >= max_col:
                col = 0
                row += 1
                if row >= max_row:
                    row = 0
                    f.write("\\end{tikzpicture}\n\n");
                    end = True
        if not end:
            f.write("\\end{tikzpicture}\n");
        f.write("\\end{document}\n");
        f.close()