# 这个是vortex的env，用于训练和评估vortex_taylor和vortex_lamb模型
# 上游干净来流，目标是vortex流过的时序信号于无pinball情况一致
# vortes系列模型都基于d1a3o12_re100模型迁移训练
import gymnasium as gym
import numpy as np
from gymnasium import spaces
import ctypes
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 = 150
CONV_LEN = 30
MAX_STEPS = 150
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=-1, high=1, shape=(S_DIM,), dtype=DATA_TYPE
        )
        self.fifo_states = deque(maxlen=FIFO_LEN)
        self.target_states = np.empty((0, 6), dtype=DATA_TYPE)
        self.force_norm_fact = 1.0
        self.sens_norm_fact = np.ones(6, dtype=DATA_TYPE)
        self.sens_deviation = np.zeros(6, dtype=DATA_TYPE)
        self.reward_cd = 0.0
        self.reward_cl = 0.0
        self.reward_sim = 0.0
        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] = (40 * L0, (NY - 1) / 2 + 2 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        center: Tuple[float, float, float] = (40 * L0, (NY - 1) / 2 - 2 * L0, 0)
        self.flow_field.add_sensor(center, L0 / 4)
        self.flow_field.run(int(1*NX/U0), np.zeros(3, dtype=DATA_TYPE))
        self.flow_field.get_ddf()
        self.flow_field.save_ddf()
        center: Tuple[float, float, float] = (10 * L0, (NY - 1) / 2, 0)
        # self.flow_field.add_vortex(center, L0 * 2, 0.5*U0, 0, "lamb")
        self.flow_field.add_vortex(center, L0 * 2, 0.03*U0, 0, "taylor") # 这里会更改vortex类型，在模型名中体现

        for i in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.zeros(3, dtype=DATA_TYPE))
            new_state = self.flow_field.obs.copy()
            self.target_states = np.vstack((self.target_states, new_state))

        self.flow_field.restore_ddf()
        self.flow_field.apply_ddf()
        center: Tuple[float, float, float] = (30 * L0, (NY - 1) / 2, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 + 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        center: Tuple[float, float, float] = (31.3 * L0, (NY - 1) / 2 - 0.75 * L0, 0)
        self.flow_field.add_cylinder(center, L0 / 2)
        self.flow_field.run(int(1*NX/U0), np.zeros(6, dtype=DATA_TYPE))
        self.flow_field.run(int(1*NX/U0), np.array([0.0, 0.0, 0.0, 0.0, -4*U0, 4*U0], dtype=DATA_TYPE))
        # self.flow_field.get_ddf()
        # self.flow_field.save_ddf()
        center: Tuple[float, float, float] = (15 * L0, (NY - 1) / 2, 0)
        # self.flow_field.add_vortex(center, L0 * 2, 0.5*U0, 0, "lamb")
        self.flow_field.add_vortex(center, L0 * 2, 0.03*U0, 0, "taylor")
        self.flow_field.save_ddf()

        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.apply_ddf()

        temp_states = np.array(self.fifo_states)
        self.force_norm_fact = 6 * np.max(np.abs(temp_states[:, 6:12]))
        for i in range(6):
            self.sens_deviation[i] = np.mean(temp_states[:, i])
            self.sens_norm_fact[i] = 5 * np.max(np.abs(temp_states[:, i] - self.sens_deviation[i]))

        for i in range(FIFO_LEN):
            self.flow_field.run(SAMPLE_INTERVAL, np.array([0.0, 0.0, 0.0, 0.0, -4*U0, 4*U0], dtype=DATA_TYPE))
            self.fifo_states.append(self.flow_field.obs.copy())

        self.save_states = self.fifo_states.copy()
        self.flow_field.apply_ddf()


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

        # barrier = threading.Barrier(2)
        result_queue = queue.Queue()

        def run_flow_field(action):
            self.flow_field.context.push()
            U0 = config_field.velocity
            try:
                temp = np.zeros(6, dtype=DATA_TYPE)
                temp[3:6] = np.array((action*4+[0,-4,4])*U0, dtype=DATA_TYPE)
                self.flow_field.run(SAMPLE_INTERVAL, temp)
            finally:
                self.flow_field.context.pop()
            # barrier.wait()
            self.fifo_states.append(self.flow_field.obs.copy())

        def proc_data():
            states = np.array(self.fifo_states)
            forces = states[-1, 6:12] / self.force_norm_fact
            cd = (forces[0] + forces[2] + forces[4]) / 3
            cl = (forces[1] + forces[3] + forces[5]) / 3
            sens = (states[-1, 0:6] - self.sens_deviation) / self.sens_norm_fact

            similarities = 0.0
                
            def calc_sim(target, state):

                n = len(target)
                m = len(state)
                
                dtw_matrix = np.full((n + 1, m + 1), np.inf)
                dtw_matrix[0, 0] = 0
                
                for i in range(1, n + 1):
                    for j in range(1, m + 1):
                        cost = abs(target[i - 1] - state[j - 1])
                        last_min = min(dtw_matrix[i - 1, j],
                                    dtw_matrix[i, j - 1],
                                    dtw_matrix[i - 1, j - 1])
                        dtw_matrix[i, j] = cost + last_min

                return 1 - (dtw_matrix[n, m] / len(target))

            for i in range(0, 6):
                target_seq = np.roll(self.target_states[-CONV_LEN:, i], -self.current_step-1)
                state_seq = states[-CONV_LEN:, i]
                similarities += calc_sim(target_seq, state_seq) / 6

            self.reward_cd = np.exp(-np.abs(cd * 20))
            self.reward_cl = np.exp(-np.abs(cl * 80))
            self.reward_sim = np.exp(-10*np.abs(similarities - 1))
            reward = np.minimum(0.2 * self.reward_cd + 0.3 * self.reward_cl + 0.5 * self.reward_sim, 1.0)
            # barrier.wait()
            result_queue.put((np.hstack([forces, sens]), reward))

        run_flow_field(action)
        proc_data()
        observation, reward = result_queue.get()

        truncated = bool(np.any(observation > 1) or np.any(observation < -1))
        observation = np.clip(observation, -1, 1)
        self.current_step += 1
        done = self.current_step >= MAX_STEPS
        return observation, float(reward), done, truncated, {}
    
    def reset(self, seed=None):
        self.flow_field.restore_ddf()
        self.flow_field.apply_ddf()
        self.fifo_states = self.save_states.copy()
        self.current_step = 0
        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__()