import gymnasium as gym
import numpy as np
from gymnasium import spaces
import ctypes
from collections import deque

lbm = ctypes.cdll.LoadLibrary('./lbm_sens.so')

S_DIM, A_DIM = 6, 3
action_amp = 5
action_weight = 0.5
sample_interval = 200
max_steps = 320

class CustomEnv(gym.Env):
    """Custom Environment that follows gym interface."""

    metadata = {"render_modes": ["human"], "render_fps": 1000/sample_interval}

    def __init__(self, devicenum=0, Ccost=0.2):
        super().__init__()
        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
        self.observation_space = spaces.Box(low=-5, high=5, shape=(6,), dtype=np.float32)
        self.fifo_rewards = deque(maxlen=50)
        lbm.SetDevice(devicenum)
        lbm.InitAll()
        lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
        lbm.CoreSolver.restype  = ctypes.POINTER(ctypes.c_float)
        self.temps_init = lbm.CoreSolver(100*1000, 0.0, 0.0, 0.0, 0.0)
        self.s = np.array([0.0] * S_DIM, dtype=np.float32)
        for i in range(S_DIM):
            self.s[i] = self.temps_init[i]
        lbm.InitCPUMemory()
        self.max_steps = max_steps
        self.current_step = 0
        self.Ccost = Ccost

    def step(self, action):
        assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
        lbm.CoreSolver.argtypes = (ctypes.c_int,ctypes.c_float,ctypes.c_float,ctypes.c_float,ctypes.c_float)
        lbm.CoreSolver.restype  = ctypes.POINTER(ctypes.c_float)
        action = action_amp * action
        temps = lbm.CoreSolver(sample_interval, action_weight, action[0], action[1], action[2])
        for i in range(S_DIM):
            self.s[i] = temps[i]
        observation = np.hstack(self.s)
        cd = self.s[0]+self.s[2]+self.s[4]
        cl = self.s[1]+self.s[3]+self.s[5]
        reward = float((1-self.Ccost)*np.exp(-np.abs(cd)/3)+self.Ccost*np.exp(-np.abs(cl)/3))
        self.fifo_rewards.append(reward)
        terminated = bool(np.mean(self.fifo_rewards) > 0.9)
        truncated = bool(np.any(self.s > 3) or np.any(self.s < -3))
        self.current_step += 1
        if self.current_step >= self.max_steps:
            terminated = True
        info = {}
        return observation, reward, terminated, truncated, info

    def reset(self, seed=None, Ccost=0.2):
        lbm.ResetAll()
        for i in range(S_DIM):
            self.s[i] = self.temps_init[i]
        observation = np.hstack(self.s)
        info = {}
        self.current_step = 0
        self.Ccost = Ccost
        return observation, info

    def render(self, episode=0, numstep=0):
        lbm.OutputFlow.argtypes = (ctypes.c_int, ctypes.c_int, ctypes.c_int)
        lbm.OutputFlow(episode, numstep, sample_interval)

    def close(self):
        lbm.Finalize()