import numpy as np
from pyopu.tensors.base_tensor import Rank1Tensor, Rank2Tensor, Rank3Tensor
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, organoid, interface):
        super().__init__()
        self.organoid = organoid
        self.interface = interface
        self.T1 = None; self.T2 = None; self.T3 = None

    def set_tensors(self, t1: Rank1Tensor, t2: Rank2Tensor = None, t3: Rank3Tensor = None):
        nv = t1.num_vars
        if t2 is not None and t2.num_vars != nv: raise ValueError("Rank-2 var mismatch.")
        if t3 is not None and t3.num_vars != nv: raise ValueError("Rank-3 var mismatch.")
        self.T1 = t1.data
        self.T2 = t2.data if t2 is not None else None
        self.T3 = t3.data if t3 is not None else None

    def _calculate_force(self):
        Force = np.copy(self.T1)
        if self.T2 is not None:
            Force += np.dot(self.T2, self.organoid.x_state) + np.dot(self.T2.T, self.organoid.x_state)
        if self.T3 is not None:
            F3 = np.einsum('ijk,j,k->i', self.T3, self.organoid.x_state, self.organoid.x_state)
            F3 += np.einsum('jik,j,k->i', self.T3, self.organoid.x_state, self.organoid.x_state)
            F3 += np.einsum('jki,j,k->i', self.T3, self.organoid.x_state, self.organoid.x_state)
            Force += F3
        return Force

    def run(self, duration_ms=1000.0, dt_kernel_ms=10.0):
        if self.T1 is None: raise ValueError("Tensors not set.")
        steps = int(duration_ms / dt_kernel_ms)
        print(f"Starting OPU Simulation ({self.organoid.num_neobits} Neobits)...")

        for step in range(steps):
            x_inst = self.interface.read_from_organoid(self.organoid._control_spike_mon, self.organoid.neurons_per_neobit, self.organoid.num_neobits)
            self.organoid.x_state = 0.8 * self.organoid.x_state + 0.2 * x_inst
            force = self._calculate_force()
            self.interface.write_to_organoid(self.organoid.neuron_group, force, self.organoid.neurons_per_neobit)
            
            # Notify observers (like EnergyTelemetry)
            self.notify_telemetries(x_state=self.organoid.x_state, force=force)
            
            self.organoid.run_step(dt_kernel_ms)
            
        print("Execution Complete.")
        return OPUResult(self.organoid.x_state)
