import json
import numpy as np
import copy
import time

class AlgebraicEquation:
    def __init__(self):
        self.strength = 1.0
        self.usage_count = 0
        self.last_used_step = 0

    def fit(self, x, y):
        pass

    def update_parameters(self, x, y):
        pass

    def predict(self, x):
        pass

    def get_error(self, x, y):
        prediction = self.predict(x)
        return abs(prediction - y) if prediction is not None else float('inf')

    def copy(self):
        return copy.deepcopy(self)

class LinearEquation(AlgebraicEquation):
    def __init__(self):
        super().__init__()
        self.c = 0.0
    def fit(self, x, y): self.c = y - x
    def update_parameters(self, x, y): self.c = y - x
    def predict(self, x): return x + self.c

class MultiplicativeEquation(AlgebraicEquation):
    def __init__(self):
        super().__init__()
        self.a = 1.0
    def fit(self, x, y): self.a = y / x if x != 0 else 0.0
    def update_parameters(self, x, y): self.a = y / x if x != 0 else 0.0
    def predict(self, x): return self.a * x

class QuadraticEquation(AlgebraicEquation):
    def __init__(self):
        super().__init__()
        self.c = 0.0
    def fit(self, x, y): self.c = y - (x ** 2)
    def update_parameters(self, x, y): self.c = y - (x ** 2)
    def predict(self, x): return (x ** 2) + self.c

class NextTokenSystem:
    def __init__(self, vocabulary):
        self.vocabulary = vocabulary
        self.token_to_score = {token: float(i + 1) for i, token in enumerate(vocabulary)}
        self.score_to_token = {score: token for token, score in self.token_to_score.items()}
        self.relationship_table = []
        self.candidate_equations = []
        self.timestep = 0

    def get_score(self, token): return self.token_to_score.get(token)

    def add_relationship(self, token_x, token_y):
        sx, sy = self.get_score(token_x), self.get_score(token_y)
        if sx is not None and sy is not None: self.relationship_table.append((sx, sy))

    def generate_candidate_rules(self):
        self.candidate_equations = []
        for x, y in self.relationship_table:
            for Cls in [LinearEquation, MultiplicativeEquation, QuadraticEquation]:
                eq = Cls(); eq.fit(x, y); self.candidate_equations.append(eq)

class ConflictResolutionEngine(NextTokenSystem):
    def __init__(self, vocabulary, decay_factor=0.9, reuse_bonus=0.2, inhibitory_penalty=10.0):
        super().__init__(vocabulary)
        self.decay_factor, self.reuse_bonus, self.inhibitory_penalty = decay_factor, reuse_bonus, inhibitory_penalty
        self.current_step, self.last_selected_eq, self.local_anchors = 0, None, {}
        self.reinforcement_bonus = 5.0

    def calculate_dynamic_threshold(self, percentage=0.0001): return max(1.0, len(self.vocabulary) * percentage)

    def register_local_anchor(self, token, equation): self.local_anchors[token] = equation.copy()

    def refined_back_search(self, input_score, target_score, threshold=None):
        if threshold is None: threshold = self.calculate_dynamic_threshold()
        best_eq, min_diff = None, float('inf')
        for eq in self.candidate_equations:
            diff = abs(eq.predict(input_score) - target_score)
            if diff <= threshold and diff < min_diff:
                min_diff, best_eq = diff, eq
        return best_eq

    def select_governing_equation(self, current_token=None):
        if current_token in self.local_anchors: return self.local_anchors[current_token]
        self.current_step += 1
        weights = []
        for i, eq in enumerate(self.candidate_equations):
            if eq.last_used_step < self.current_step - 1: eq.strength *= self.decay_factor
            w = eq.strength + (self.relationship_table[i//3][0] + self.relationship_table[i//3][1]) / (2 * len(self.vocabulary))
            weights.append(max(0, w))
        best_eq = self.candidate_equations[np.argmax(weights)]
        self.last_selected_eq = best_eq
        return best_eq

    def predict_next_token(self, current_token):
        score = self.get_score(current_token)
        eq = self.select_governing_equation(current_token)
        target_y = eq.predict(score)
        best_token = self.score_to_token[min(self.score_to_token.keys(), key=lambda k: abs(k - target_y))]
        return {'predicted_token': best_token, 'governing_equation': type(eq).__name__}

    def penalized_supervised_train(self, data_path, iterations=25, threshold=None):
        with open(data_path, 'r') as f: examples = json.load(f)['training_examples']
        if threshold is None: threshold = self.calculate_dynamic_threshold()
        for i in range(iterations):
            for ex in examples:
                res = self.predict_next_token(ex['input'])
                if res['predicted_token'] != ex['target']:
                    self.last_selected_eq.strength = max(0, self.last_selected_eq.strength - self.inhibitory_penalty)
                    best = self.refined_back_search(self.get_score(ex['input']), self.get_score(ex['target']), threshold)
                    if best:
                        opt = best.copy(); opt.update_parameters(self.get_score(ex['input']), self.get_score(ex['target']))
                        self.register_local_anchor(ex['input'], opt)

print('model.py has been successfully written to the current directory.')