from jinja2 import Environment, FileSystemLoader
from src.models import OuterKey, InnerKey, SubstitutionKey


def render_spn_algebra(data: dict):
    # Set up the Jinja2 environment and template loader
    file_loader = FileSystemLoader('./docs')
    env = Environment(loader=file_loader)

    # Load the template
    template = env.get_template('spn_algebra_template.md')
    # Render the template with the data
    output = template.render(data)

    # Print or save the output
    output_file = "/Users/donov/Desktop/NKode_documentation/nkode/docs/spn_algebra.md"
    with open(output_file, 'w') as fp:
        fp.write(output)
        print("File written successfully")


if __name__ == "__main__":
    height = 9
    width = 8
    total_attr = height * width
    i0 = InnerKey.init_matrix(width, height)
    i1 = InnerKey.init_matrix(width, height)
    i2 = InnerKey.init_matrix(width, height)
    i_identity = InnerKey.init_identity_matrix(width, height)

    o0 = OuterKey.init_matrix(height)
    o1 = OuterKey.init_matrix(height)
    o2 = OuterKey.init_matrix(height)
    o_identity = OuterKey.init_identity_matrix(height)

    s0 = SubstitutionKey.init_matrix(width, height)
    s1 = SubstitutionKey.init_matrix(width, height)
    s2 = SubstitutionKey.init_matrix(width, height)
    s_identity = SubstitutionKey.init_identity_matrix(width, height)

    s0_o0 = s0 < o0

    i0_o0 = i0 < o0

    o0_o1 = o0 << o1

    i0_i1 = i0 << i1

    s0_i0 = s0 << i0

    inv_o0 = ~o0

    inv_i0 = ~i0

    s0_s1 = s0 ^ s1

    data = {
        'height': height,
        'width': width,
        'total_attr': total_attr,
        'i0': i0,
        'i1': i1,
        'i2': i2,
        'i_identity': i_identity,
        'o0': o0,
        'o1': o1,
        'o2': o2,
        'o_identity': o_identity,
        's0': s0,
        's1': s1,
        's2': s2,
        's_identity': s_identity,
        's0_o0': s0_o0,
        'i0_o0': i0_o0,
        'o0_o1': o0_o1,
        'i0_i1': i0_i1,
        's0_i0': s0_i0,
        'inv_o0': inv_o0,
        'inv_i0': inv_i0,
        's0_s1': s0_s1
    }

    render_spn_algebra(data)