import gymnasium as gym
import numpy as np
from gymnasium import spaces
from collections import deque
from typing import Tuple
import sys
import os
import matplotlib.pyplot as plt
import queue

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"

current_dir = os.path.dirname(os.path.abspath("__file__"))
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(parent_dir)
from CelerisLab import FlowField
from CelerisLab import utils

config_cuda = utils.load_cuda_config(
    os.path.join(parent_dir, "configs", "config_cuda.json")
)
config_field = utils.load_flow_field_config(
    os.path.join(parent_dir, "configs", "config_flowfield.json")
)

S_DIM, A_DIM = 12, 3
U0 = config_field.velocity
T0 = 1000
SAMPLE_INTERVAL = 800
FIFO_LEN = 360
MAX_STEPS = 500
if config_field.data_type == "FP32":
    DATA_TYPE = np.float32
else:
    raise ValueError(f"Unsupported data type {config_field.data_type}.")


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

    metadata = {"render_modes": ["human"], "render_fps": T0 / SAMPLE_INTERVAL}

    def __init__(self, device_id=0):
        super().__init__()
        self.action_space = spaces.Box(low=-1, high=1, shape=(A_DIM,), dtype=DATA_TYPE)
        self.observation_space = spaces.Box(
            low=-5, high=5, shape=(S_DIM,), dtype=DATA_TYPE
        )
        self.fifo_states = deque(maxlen=FIFO_LEN)
        self.save_states = deque(maxlen=FIFO_LEN)
        self.current_step = 0

        self.flow_field = FlowField(config_field, config_cuda, device_id)
        L0 = 20
        U0 = config_field.velocity
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        center: Tuple[float, float, float] = (20 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (21.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (21.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 + 3 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2 - 3 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        self.flow_field.run(int(4*NX/U0), np.zeros(6, dtype=DATA_TYPE))

        for i in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(6, dtype=DATA_TYPE))
            self.fifo_states.append(self.flow_field.obs.copy())
        
        self.save_states = self.fifo_states.copy()
        self.flow_field.get_ddf()
        self.flow_field.save_ddf()
            

    def step(self, action):
        assert self.action_space.contains(action), "%r (%s) invalid" % (
            action,
            type(action),
        )

        def run_flow_field(action):
            self.flow_field.context.push()
            U0 = config_field.velocity
            NX = self.flow_field.FIELD_SHAPE[0]
            L0 = 20
            try:
                temp = np.zeros(6, dtype=DATA_TYPE)
                temp[0:3] = np.array((action*4+[0,-4,+4])*U0, dtype=DATA_TYPE)                
                if self.current_step == 0:
                    self.flow_field.run(int(2*NX/U0), temp)
                self.flow_field.run(SAMPLE_INTERVAL, temp)
            finally:
                state = self.flow_field.obs.copy()
                force = state[0:6] / (L0*U0*U0)
                sens = state[6:12] / 78 / U0
                self.fifo_states.append([force, sens])
                self.flow_field.context.pop()

        run_flow_field(action)

        truncated = False
        observation = 0.0
        terminated = self.current_step >= MAX_STEPS
        self.current_step += 1
        return observation, 0.0, terminated, truncated, {}
    
    def reset(self, seed=None):
        self.flow_field.apply_ddf()
        self.current_step = 0
        self.fifo_states = self.save_states.copy()
        return np.zeros(S_DIM, dtype=np.float32), {}
    
    def render(self, mode="human"):
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        self.flow_field.get_ddf()
        ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
        ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
        uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
        speed = np.sqrt(ux**2 + uy**2)
        plt.figure(figsize=(10, 5))
        plt.imshow(speed.T, origin='lower', cmap='viridis', extent=[0, NX, 0, NY])
        plt.colorbar(label='Speed')
        plt.title('Scalar Velocity Field')
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.tight_layout()
        plt.show()

    def save_field(self, filename):
        NX = self.flow_field.FIELD_SHAPE[0]
        NY = self.flow_field.FIELD_SHAPE[1]
        self.flow_field.get_ddf()
        ddf_plot = self.flow_field.ddf.copy().reshape((9, NY, NX)).transpose(2, 1, 0)
        flag_plot = self.flow_field.flag.copy().reshape((NY, NX)).transpose(1, 0)
        ux = (ddf_plot[:, :, 1] + ddf_plot[:, :, 5] + ddf_plot[:, :, 8] - ddf_plot[:, :, 3] - ddf_plot[:, :, 6] - ddf_plot[:, :, 7]) / U0
        uy = (ddf_plot[:, :, 2] + ddf_plot[:, :, 5] + ddf_plot[:, :, 6] - ddf_plot[:, :, 4] - ddf_plot[:, :, 7] - ddf_plot[:, :, 8]) / U0
        with open(os.path.join(parent_dir, "output", filename), "w") as f:
            f.write("Title= \"LBM 2D\"\r\n")
            f.write("VARIABLES= \"X\",\"Y\",\"flag\",\"U\",\"V\",\r\n")
            f.write(f"ZONE T= \"BOX\",I= {NX},J= {NY},F=POINT\r\n")
            for j in range(NY):
                for i in range(NX):
                    f.write(f"{i},{j},{flag_plot[i, j]},{ux[i, j]},{uy[i, j]}\r\n")

    def close(self):
        self.flow_field.__del__()