# Filename: star_complement.py
import numpy as np
from itertools import product, permutations, combinations
from numpy.linalg import LinAlgError, eigvalsh
from collections import defaultdict

def high_precision_spectrum(matrix, tol=1e-6):
    eigenvalues = eigvalsh(matrix)
    rounded = np.round(eigenvalues, decimals=10)
    return np.sort(rounded)

def bron_kerbosch_pivot(R, P, X, neighbors, cliques):
    if not P and not X:
        cliques.append(R.copy())
        return
    u = max(P.union(X), key=lambda vertex: len(neighbors[vertex]), default=None)
    for v in P - neighbors.get(u, set()):
        bron_kerbosch_pivot(
            R.union({v}),
            P.intersection(neighbors[v]),
            X.intersection(neighbors[v]),
            neighbors,
            cliques
        )
        P.remove(v)
        X.add(v)

def find_maximal_cliques(compatibility_graph):
    vertices = set(compatibility_graph.keys())
    cliques = []
    bron_kerbosch_pivot(set(), vertices, set(), compatibility_graph, cliques)
    return cliques

def are_switching_isomorphic(A, B, tol=1e-6):
    n = A.shape[0]
    if B.shape != (n, n):
        return False

    deg_A = np.sum(A != 0, axis=1)
    deg_B = np.sum(B != 0, axis=1)
    if not np.array_equal(np.sort(deg_A), np.sort(deg_B)):
        return False

    spec_A = high_precision_spectrum(A, tol)
    spec_B = high_precision_spectrum(B, tol)
    if not np.allclose(spec_A, spec_B, atol=tol):
        return False

    for pi in permutations(range(n)):
        P = np.eye(n)[list(pi)]
        A_perm = P @ A @ P.T
        for signs in product([-1, 1], repeat=n):
            S = np.diag(signs)
            A_test = S @ A_perm @ S
            if np.allclose(A_test, B, atol=tol):
                return True
    return False

def remove_redundant_vertices(A_ext, k, lam):
    n = A_ext.shape[0]
    ext_indices = set(range(k, n))
    to_remove = set()

    for u, v in combinations(range(n), 2):
        edge = A_ext[u, v]
        if lam == 0 and edge != 0:
            continue
        if lam == 1 and edge != -1:
            continue
        if lam == -1 and edge != 1:
            continue

        nu_u = np.delete(A_ext[u], [u, v])
        nu_v = np.delete(A_ext[v], [u, v])
        if np.array_equal(nu_u, nu_v) or np.array_equal(nu_u, -nu_v):
            if u in ext_indices and v not in ext_indices:
                to_remove.add(u)
            elif v in ext_indices and u not in ext_indices:
                to_remove.add(v)

    if not to_remove:
        return A_ext

    keep = [i for i in range(n) if i not in to_remove]
    return A_ext[np.ix_(keep, keep)]

def read_input_file(filename="input.txt"):
    with open(filename, 'r') as f:
        lines = f.readlines()

    matrix_lines = [line.strip() for line in lines if line.strip() and not line.startswith("lambda")]
    A = np.array([list(map(float, line.split())) for line in matrix_lines])

    lambda_line = [line for line in lines if line.strip().startswith("lambda")]
    if not lambda_line:
        raise ValueError("No lambda value found in input file.")
    lam = float(lambda_line[0].split('=')[1])

    return A, lam

def write_output_file(solutions, filename="output.txt"):
    with open(filename, 'w') as f:
        for idx, A_ext in enumerate(solutions):
            f.write(f"Solution {idx+1}:\n")
            for row in A_ext:
                f.write(' '.join(f"{val:.0f}" for val in row) + '\n')
            spectrum = high_precision_spectrum(A_ext)
            f.write("Spectrum: " + ', '.join(map(str, spectrum)) + '\n\n')
    print(f"✓ Output written to {filename}")

def write_bvectors(b_vectors, filename="bvectors.txt"):
    with open(filename, 'w') as f:
        for b in b_vectors:
            f.write(' '.join(map(str, b)) + '\n')
    print(f"✓ b-vectors written to {filename}")

def write_mextensions(extensions, filename="mextensions.txt"):
    with open(filename, 'w') as f:
        for idx, A_ext in enumerate(extensions):
            f.write(f"Extension {idx+1}:\n")
            for row in A_ext:
                f.write(' '.join(f"{val:.0f}" for val in row) + '\n')
            spectrum = high_precision_spectrum(A_ext)
            f.write("Spectrum: " + ', '.join(map(str, spectrum)) + '\n\n')
    print(f"✓ Maximal extensions written to {filename}")

def write_compatibility_graph(neighbors, filename="compatibility_graph.txt"):
    nodes = sorted(neighbors.keys())
    n = len(nodes)
    index_map = {node: idx for idx, node in enumerate(nodes)}
    adj_matrix = np.zeros((n, n), dtype=int)
    for u in nodes:
        for v in neighbors[u]:
            i, j = index_map[u], index_map[v]
            adj_matrix[i][j] = 1
    with open(filename, 'w') as f:
        for row in adj_matrix:
            f.write(' '.join(map(str, row)) + '\n')
    print(f"✓ Compatibility graph (adjacency matrix) written to {filename}")

def find_star_complement_extensions(A_H, lam, tol=1e-6):
    k = A_H.shape[0]
    spectrum_H = high_precision_spectrum(A_H, tol)
    if np.any(np.isclose(spectrum_H, lam, atol=tol)):
        print(f"λ = {lam} is in the spectrum of the initial graph. No extensions are computed.")
        return []

    try:
        M_inv = np.linalg.inv(lam * np.eye(k) - A_H)
    except LinAlgError:
        print("Matrix (λ I - A_H) is singular.")
        return []

    b_vectors = []
    seen = set()
    A_H_cols = [A_H[:, i] for i in range(k)] if lam in {0, 1, -1} else []

    for vec in product([-1, 0, 1], repeat=k):
        if all(x == 0 for x in vec):
            continue
        key = tuple(vec)
        neg_key = tuple(-x for x in vec)
        if neg_key in seen:
            continue
        b = np.array(vec, dtype=float)
        quad = b.T @ M_inv @ b
        if np.isclose(quad, lam, atol=tol):
            b_vectors.append(b)
            seen.add(key)

    print(f"✓ Found {len(b_vectors)} candidate b-vectors.")
    write_bvectors(b_vectors)

    edge_labels = {}
    neighbors = defaultdict(set)

    for i, j in combinations(range(len(b_vectors)), 2):
        bi = b_vectors[i]
        bj = b_vectors[j]
        valid_labels = []
        for label, sign in zip(['0', '+', '-'], [0, 1, -1]):
            A_S = np.array([[0, sign], [sign, 0]])
            B = np.vstack([bi, bj])
            A_ext = np.block([
                [A_H, B.T],
                [B, A_S]
            ])
            spec = high_precision_spectrum(A_ext)
            count = np.count_nonzero(np.isclose(spec, lam, atol=tol))
            if count == 2:
                valid_labels.append(label)
        if len(valid_labels) == 1:
            label = valid_labels[0]
            edge_labels[(i, j)] = label
            edge_labels[(j, i)] = label
            neighbors[i].add(j)
            neighbors[j].add(i)

    write_compatibility_graph(neighbors)
    cliques = find_maximal_cliques(neighbors)
    print(f"✓ Found {len(cliques)} maximal cliques.")

    all_extensions = []
    for clique in cliques:
        idx = list(clique)
        if not idx:
            continue
        B = np.vstack([b_vectors[i] for i in idx])
        s = len(idx)
        A_S = np.zeros((s, s))
        for i in range(s):
            for j in range(i+1, s):
                label = edge_labels.get((idx[i], idx[j]), '0')
                if label == '+':
                    A_S[i, j] = A_S[j, i] = 1
                elif label == '-':
                    A_S[i, j] = A_S[j, i] = -1
        A_ext = np.block([
            [A_H, B.T],
            [B, A_S]
        ])
        if lam in {0, 1, -1}:
            A_ext = remove_redundant_vertices(A_ext, k, lam)
        all_extensions.append(A_ext)

    write_mextensions(all_extensions)

    unique_extensions = []
    extensions_seen = []
    for A_ext in all_extensions:
        is_new = True
        for existing in extensions_seen:
            if are_switching_isomorphic(A_ext, existing, tol):
                is_new = False
                break
        if is_new:
            unique_extensions.append(A_ext)
            extensions_seen.append(A_ext)

    print(f"✓ Found {len(unique_extensions)} switching non-isomorphic extensions.")
    return unique_extensions

if __name__ == "__main__":
    try:
        A_H, lam = read_input_file("input.txt")
        results = find_star_complement_extensions(A_H, lam)
        if results:
            write_output_file(results, "output.txt")
    except Exception as e:
        print("Error:", e)