import numpy as np
import itertools
import matplotlib.pyplot as plt

class NeoTerm:
    def __init__(self, weight, logic):
        self.weight = weight
        self.logic = logic

    def get_indices_and_signs(self):
        indices = []
        signs = []
        for var in self.logic:
            sign = not var.startswith("~")
            idx = int(var.replace("~", "").replace("x", ""))
            indices.append(idx)
            signs.append(sign)
        return indices, signs

class NeoCode:
    def __init__(self, num_neobits):
        self.num_neobits = num_neobits
        self.terms = []

    def add_term(self, term):
        self.terms.append(term)

    def draw(self, figsize=(6, 3)):
        fig, ax = plt.subplots(figsize=figsize)
        
        # Horizontal lines for neobits
        for i in range(self.num_neobits):
            ax.hlines(y=i, xmin=-0.5, xmax=len(self.terms) - 0.5, color='navy', zorder=1)
            ax.text(-0.6, i, f'$x_{{{i}}}$', va='center', ha='right', fontsize=12)

        eq_terms = []
        for j, term in enumerate(self.terms):
            indices, signs = term.get_indices_and_signs()
            if not indices:
                # Rank 0 term
                term_str = str(term.weight) if term.weight < 0 else f"+{term.weight}"
                eq_terms.append(term_str)
                continue
            
            # Vertical line connecting the highest and lowest index in this term
            min_idx, max_idx = min(indices), max(indices)
            ax.vlines(x=j, ymin=min_idx, ymax=max_idx, color='gray', linestyle='--', zorder=2)
            
            term_str_parts = []
            for idx, sign in zip(indices, signs):
                if sign:
                    ax.plot(j, idx, 'ko', markersize=8, zorder=3)
                    term_str_parts.append(f'x_{{{idx}}}')
                else:
                    ax.plot(j, idx, 'wo', markeredgecolor='k', markersize=8, markeredgewidth=2, zorder=3)
                    ax.plot(j, idx, 'ko', markersize=3, zorder=3) 
                    term_str_parts.append(f'(1 - x_{{{idx}}})')
            
            # Equation for this term
            term_str = " ".join(term_str_parts)
            term_expr = f"{term.weight}{term_str}" if term.weight < 0 else f"+{term.weight}{term_str}"
            eq_terms.append(term_expr)

        ax.set_yticks(range(self.num_neobits))
        ax.set_yticklabels([])
        ax.set_xticks(range(len(self.terms)))
        ax.set_xticklabels([f'T{j}' for j in range(len(self.terms))])
        ax.invert_yaxis()
        
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.tick_params(axis='both', which='both', length=0)
        
        equation = "$H(\\mathbf{x}) = " + " ".join(eq_terms) + "$"
        equation = equation.replace("= +", "= ")
        
        plt.figtext(0.5, 0.01, equation, ha="center", fontsize=14)
        plt.tight_layout(rect=[0, 0.08, 1, 1])
        plt.show()

class Compiler:
    def __init__(self, num_neobits=None):
        self.num_neobits = num_neobits

    def compile(self, neocode):
        if self.num_neobits is None:
            num_neobits = neocode.num_neobits
        else:
            num_neobits = self.num_neobits

        tensors = {}
        for term in neocode.terms:
            indices, signs = term.get_indices_and_signs()
            
            P = [i for i, s in zip(indices, signs) if s]
            N = [i for i, s in zip(indices, signs) if not s]
            
            for r in range(len(N) + 1):
                for S in itertools.combinations(N, r):
                    active_vars = tuple(sorted(P + list(S)))
                    rank = len(active_vars)
                    coeff = term.weight * ((-1) ** len(S))
                    
                    if rank not in tensors:
                        if rank == 0:
                            tensors[rank] = 0.0
                        else:
                            tensors[rank] = np.zeros([num_neobits] * rank)
                        
                    if rank == 0:
                        tensors[rank] += coeff
                    else:
                        perms = list(itertools.permutations(active_vars))
                        unique_perms = list(set(perms))
                        sym_coeff = coeff / len(unique_perms)
                        
                        for perm in unique_perms:
                            tensors[rank][perm] += sym_coeff
                        
        return tensors
