from libsolve import *
from fractions import Fraction
import numpy as np

FOUR_SIMEQ = '\\overset{{\\texttt{{two.py}}}}\\simeq'

# generate figure as plain file, so we can \input it
def gen_figure(figname, text):
    with open(f'figures/fig_four_{figname}.tex', 'w') as fd:
        fd.write('$$' + text + '$$')

A = Matrix([
    [-1, 1, 1, -1],
    [-7, 4, 4, -3],
    [4, -1, -2, 1],
    [4, -1, -2, 1]
])

lbdE = Matrix([
    [Poly([0, 1]), 0, 0, 0],
    [0, Poly([0, 1]), 0, 0],
    [0, 0, Poly([0, 1]), 0],
    [0, 0, 0, Poly([0, 1])]
])

print((A - lbdE).det())






A = A - Matrix([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, 1]
])

A = A@A

before = A.to_tex()

A.triangulate()

A.make_D(0, -1)
A.make_U(1, 2)
A.make_D(1, Fraction(1, 4))

gen_figure('V_1',
f'''
    {before}
    {FOUR_SIMEQ}
    {A.to_tex()}
''')








A = Matrix([
    [-1, 1, 1, -1],
    [-7, 4, 4, -3],
    [4, -1, -2, 1],
    [4, -1, -2, 1]
])
A = A@A

before = A.to_tex()

A.triangulate()


A.make_D(0, -Fraction(1, 2))
A.make_D(1, -6)

gen_figure('V_0',
f'''
    {before}
    {FOUR_SIMEQ}
    {A.to_tex()}
''')






# solve 16-equation system
# unknown matrix is (answer_to_system).reshape(4, 4),
# so we 'flatten' it to solve for each unknown

def make_row(vec, ans):
    return [
        [0, 0, 0, 0] * i +
        vec +
        [0, 0, 0, 0] * (3 - i) +
        [ans[i]]
        for i in range(4)
    ]

v_11 = [1, 3, 0, 0]
v_12 = [1, 0, 3, 3]
v_01 = [-1,-3, 1, 0]
v_02 = [12, 10, 0, 3]
e_3 = [0, 0, 1, 0]
e_4 = [0, 0, 0, 1]

rows = []
rows += make_row(v_11, [0, 0, 0, 0])
rows += make_row(v_12, [0, 0, 0, 0])
rows += make_row(e_3, v_01)
rows += make_row(e_4, v_02)

sys = Matrix(rows)

before = sys.to_tex()

sys.triangulate(True)
sys.make_D(1, -Fraction(1, 3))
sys.make_D(5, -Fraction(1, 3))
sys.make_D(9, -Fraction(1, 3))
sys.make_D(13, -Fraction(1, 3))

gen_figure('sys_orig',
f'''
    \\begingroup % keep the change local
    \\setlength\\arraycolsep{{4pt}}
    {before}
    {FOUR_SIMEQ}
    {sys.to_tex()}
    \\endgroup
''')

ans = [row[-1] for row in sys.rows]
reshape = [ans[4 * i:4 * i + 4] for i in range(4)]

psi = Matrix(reshape)

gen_figure('psi', 
f'''
\\psi' = {psi.to_tex()}
''')






inverse = Matrix([
    [ 0,  Fraction(1, 3), 0, 0,   1, 0, 0, 0],
    [ 1, -Fraction(1, 3), 0, 0,   0, 1, 0, 0],
    [-3, -1, 1, 0,                0, 0, 1, 0],
    [-3, -1, 0, 1,                0, 0, 0, 1]
])

before = inverse.to_tex()

inverse.triangulate(True)
inverse.make_D(1, 3)

gen_figure('inverse',
f'''
    {before}
    {FOUR_SIMEQ}
    {inverse.to_tex()}
''')




C = Matrix([
    [ 0,  Fraction(1, 3), 0, 0],
    [ 1, -Fraction(1, 3), 0, 0],
    [-3, -1, 1, 0],
    [-3, -1, 0, 1]
])

inv = Matrix([
    [1, 1, 0, 0],
    [3, 0, 0, 0],
    [6, 3, 1, 0],
    [6, 3, 0, 1]
])


gen_figure('inverse_mul',
f'''
    \\psi = C^{{-1}}\\psi'C =
    {((inv @ psi) @ C).to_tex()}
'''
)