import numpy as np
from pyopu.telemetry.base_telemetry import Observable

class OPUResult:
    def __init__(self, final_probabilities):
        self.final_probabilities = final_probabilities

    def get_binary_state(self, threshold=0.5):
        return [1 if p > threshold else 0 for p in self.final_probabilities]

class TensorEngine(Observable):
    def __init__(self):
        super().__init__()
        self.tensors = {}

    def set_tensors(self, tensors: dict):
        for k, v in tensors.items():
            if not isinstance(k, int):
                raise ValueError(f"Tensor keys must be integers (rank), got {type(k)}.")
        self.tensors = tensors

    def _calculate_force(self, x_state):
        Force = np.zeros_like(x_state)
        for rank, tensor in self.tensors.items():
            if rank == 0:
                continue
            elif rank == 1:
                Force += tensor
            elif rank == 2:
                Force += 2 * np.dot(tensor, x_state)
            else:
                letters = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
                if rank > len(letters):
                    raise ValueError(f"Rank {rank} too high for automatic einsum.")
                tensor_subscript = letters[:rank]
                x_subscripts = ','.join(letters[1:rank])
                einsum_str = f"{tensor_subscript},{x_subscripts}->{letters[0]}"
                operands = [tensor] + [x_state] * (rank - 1)
                Force += rank * np.einsum(einsum_str, *operands)
        return Force

    def run_loop(self, interface, organoid, duration_ms=1000.0, dt_kernel_ms=10.0):
        if not self.tensors: raise ValueError("Tensors not set.")
        steps = int(duration_ms / dt_kernel_ms)
        import sys
        import time
        
        # Reset telemetry accumulation and interface states
        self.reset_telemetries()
        organoid.reset_telemetries()
        interface.reset()
        
        print(f"Starting OPU Simulation ({organoid.num_neobits} Neobits)...")
        start_time = time.time()

        for step in range(steps):
            x_inst = interface.read_from_organoid(organoid._control_spike_mon, organoid.neurons_per_neobit, organoid.num_neobits)
            organoid.x_state = 0.8 * organoid.x_state + 0.2 * x_inst
            force = self._calculate_force(organoid.x_state)
            interface.write_to_organoid(organoid.neuron_group, force, organoid.neurons_per_neobit)
            
            # Notify observers (like EnergyTelemetry)
            self.notify_telemetries(x_state=organoid.x_state, force=force)
            
            organoid.run_step(dt_kernel_ms)
            
            # Progress bar
            progress = (step + 1) / steps
            bar_len = 30
            filled = int(bar_len * progress)
            bar = '█' * filled + ' ' * (bar_len - filled)
            percent = progress * 100
            elapsed = time.time() - start_time
            mins, secs = divmod(int(elapsed), 60)
            sys.stdout.write(f"\r[Elapsed Time: {mins}:{secs:02d}] |{bar}| {percent:.1f}%")
            sys.stdout.flush()
            
        print("\nExecution Complete.")
        return OPUResult(organoid.x_state)
