import asyncio
import websockets
import json
import numpy as np
from typing import List, Any, Optional, Dict

class TensorCoreArray:
    def __init__(self, num_tensor_cores: int, bits: int, bandwidth_tbps: float):
        self.num_tensor_cores = num_tensor_cores
        self.bits = bits
        self.bandwidth_tbps = bandwidth_tbps
        self.initialized = False

    def initialize(self):
        print(f"Initializing {self.num_tensor_cores} tensor cores with {self.bits}-bit precision...")
        self.initialized = True

    def matmul(self, matrix_a: List[List[float]], matrix_b: List[List[float]]) -> List[List[float]]:
        if not self.initialized:
            raise RuntimeError("Tensor cores not initialized. Call initialize() first.")

        np_a = np.array(matrix_a)
        np_b = np.array(matrix_b)

        if np_a.shape[1] != np_b.shape[0]:
            raise ValueError("Matrix dimensions incompatible for multiplication")

        result = np.matmul(np_a, np_b)
        return result.tolist()

    async def send_tensor_data(self, uri: str, tensor_id: str, data: np.ndarray):
        async with websockets.connect(uri) as websocket:
            payload = {
                "operation": "tensor_data",
                "type": "send",
                "tensor_id": tensor_id,
                "data": data.tolist()
            }
            await websocket.send(json.dumps(payload))
            response = await websocket.recv()
            return json.loads(response)

    async def receive_tensor_data(self, uri: str, tensor_id: str) -> Optional[np.ndarray]:
        async with websockets.connect(uri) as websocket:
            payload = {
                "operation": "tensor_data",
                "type": "receive",
                "tensor_id": tensor_id
            }
            await websocket.send(json.dumps(payload))
            response = await websocket.recv()
            response_data = json.loads(response)
            if response_data.get("status") == "success":
                return np.array(response_data["data"])
            return None

    def get_status(self) -> Dict[str, Any]:
        return {
            "num_tensor_cores": self.num_tensor_cores,
            "bits": self.bits,
            "bandwidth_tbps": self.bandwidth_tbps,
            "initialized": self.initialized
        }

if __name__ == "__main__":
    async def test_tensor_core_array():
        tca = TensorCoreArray(num_tensor_cores=10, bits=32, bandwidth_tbps=1.0)
        tca.initialize()

        matrix_a = [[1, 2], [3, 4]]
        matrix_b = [[5, 6], [7, 8]]

        result = tca.matmul(matrix_a, matrix_b)
        print(f"Matrix multiplication result: {result}")

        # Example of sending/receiving tensor data (requires a running WebSocket server)
        # uri = "ws://localhost:7860/ws"
        # tensor_id = "test_tensor"
        # dummy_data = np.array([[10, 20], [30, 40]])
        #
        # print(f"Sending tensor data: {dummy_data.tolist()}")
        # send_response = await tca.send_tensor_data(uri, tensor_id, dummy_data)
        # print(f"Send response: {send_response}")
        #
        # received_data = await tca.receive_tensor_data(uri, tensor_id)
        # print(f"Received tensor data: {received_data.tolist() if received_data is not None else None}")

    asyncio.run(test_tensor_core_array())